diff --git a/.gitignore b/.gitignore
index 6049c2cdbe..63408699f4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,3 +6,4 @@ dist
 .idea
 .vscode
 tmp/
+requirements-musa.txt
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 573ff399c5..e7e043a1f7 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -10,4 +10,4 @@ repos:
     rev: 6.1.0 
     hooks:
       - id: flake8
-        args: ['--max-line-length=120', '--ignore=TYP001, E722, C901, E203, E266, E402, E302, E241, E902, E731, F403, E701, F405, F401, W292, W293, W503, W606, E231']
+        args: ['--max-line-length=120', '--ignore=TYP001, E722, C901, E203, E266, E402, E302, E241, E902, E731, F403, E701, F405, F401, W292, W293, W503, W606, E231, F541']
diff --git a/docs/CN/source/getting_started/installation.rst b/docs/CN/source/getting_started/installation.rst
index fb998b7567..5fa0e304d2 100755
--- a/docs/CN/source/getting_started/installation.rst
+++ b/docs/CN/source/getting_started/installation.rst
@@ -27,7 +27,7 @@ Lightllm 是一个纯python开发的推理框架，其中的算子使用triton
     $ # 前请确保你的docker设置中已经分配了足够的共享内存，否则可能导致
     $ # 服务无法正常启动。
     $ # 1.如果是纯文本服务，建议分配2GB以上的共享内存, 如果你的内存充足，建议分配16GB以上的共享内存.
-    $ # 2.如果是多模态服务，建议分配16GB以上的共享内存，具体可以根据实际情况进行调整. 
+    $ # 2.如果是多模态服务，建议分配16GB以上的共享内存，具体可以根据实际情况进行调整.
     $ # 如果你没有足够的共享内存，可以尝试在启动服务的时候调低 --running_max_req_size 参数，这会降低
     $ # 服务的并发请求数量，但可以减少共享内存的占用。如果是多模态服务，也可以通过降低 --cache_capacity
     $ # 参数来减少共享内存的占用。
@@ -38,7 +38,7 @@ Lightllm 是一个纯python开发的推理框架，其中的算子使用triton
 你也可以使用源码手动构建镜像并运行,建议手动构建镜像,因为更新比较频繁：
 
 .. code-block:: console
-    
+
     $ # 进入代码仓库的根目录
     $ cd /lightllm
     $ # 手动构建镜像, docker 目录下有不同功能场景的镜像构建文件，按需构建。
@@ -52,7 +52,7 @@ Lightllm 是一个纯python开发的推理框架，其中的算子使用triton
 或者你也可以直接使用脚本一键启动镜像并且运行：
 
 .. code-block:: console
-    
+
     $ # 查看脚本参数
     $ python tools/quick_launch_docker.py --help
 
@@ -80,6 +80,10 @@ Lightllm 是一个纯python开发的推理框架，其中的算子使用triton
     $ # 安装lightllm的依赖 (cuda 12.4)
     $ pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu124
     $
+    $ # 安装lightllm的依赖 (摩尔线程 GPU)
+    $ ./generate_requirements_musa.sh
+    $ pip install -r requirements-musa.txt
+    $
     $ # 安装lightllm
     $ python setup.py install
 
@@ -97,6 +101,6 @@ Lightllm 是一个纯python开发的推理框架，其中的算子使用triton
     .. code-block:: console
 
         $ pip install -U --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/ triton-nightly --no-deps
-    
+
     具体原因可以参考：`issue <https://github.com/triton-lang/triton/issues/3619>`_ 和 `fix PR <https://github.com/triton-lang/triton/pull/3638>`_
 
diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md
index 49b47ffa26..5d34cf747e 100755
--- a/docs/CN/source/models/add_new_model.md
+++ b/docs/CN/source/models/add_new_model.md
@@ -162,19 +162,6 @@ class BloomPreAndPostLayerWeight(PreAndPostLayerWeight):
                                                                  self.tp_rank_: split_vob_size * (self.tp_rank_ + 1), :])
             self.lm_head_weight_ = self.wte_weight_
         return
-    
-    def verify_load(self):
-        errors = "weights load not ok"
-        weights = [self.pre_norm_weight_, 
-                   self.pre_norm_bias_, 
-                   self.final_norm_weight_, 
-                   self.final_norm_bias_,
-                   self.wte_weight_,
-                   self.lm_head_weight_]
-        for i in range(len(weights)):
-            assert weights[i] is not None, "index:" + str(i) + " " + errors
-        return 
-
 ~~~
 
 ***transformer_layer_weight.py***
@@ -204,30 +191,6 @@ class BloomTransformerLayerWeight(TransformerLayerWeight):
         self._load_qkvo_weights(weights)
         self._load_ffn_weights(weights)
         return
-    
-    def verify_load(self):
-        errors = "weights load not ok"
-        weights = [self.att_norm_weight_,
-                   self.att_norm_bias_,
-                   self.q_weight_,
-                   self.k_weight_,
-                   self.v_weight_,
-                   self.q_bias_,
-                   self.k_bias_,
-                   self.v_bias_,
-                   self.o_weight_,
-                   self.o_bias_,
-
-                   self.ffn_norm_weight_,
-                   self.ffn_norm_bias_,
-                   self.ffn_1_weight_,
-                   self.ffn_1_bias_,
-                   self.ffn_2_weight_,
-                   self.ffn_2_bias_,
-                   ]
-        for i in range(len(weights)):
-            assert weights[i] is not None, "index:" + str(i) + " " + errors
-        return 
 
     def _load_qkvo_weights(self, weights):
         if f"h.{self.layer_num_}.input_layernorm.weight" in weights:
diff --git a/docs/CN/source/tutorial/api_server_args_zh.rst b/docs/CN/source/tutorial/api_server_args_zh.rst
index ce7a79ab97..3503ea9ca4 100755
--- a/docs/CN/source/tutorial/api_server_args_zh.rst
+++ b/docs/CN/source/tutorial/api_server_args_zh.rst
@@ -183,22 +183,6 @@ PD 分离模式参数
     设置为 True 时，--nccl_host 必须等于 config_server_host，--nccl_port 对于 config_server 必须是唯一的，
     不要为不同的推理节点使用相同的 nccl_port，这将是严重错误
 
-attention类型选择参数
----------------------
-
-.. option:: --mode
-
-    模型推理模式，可以指定多个值：
-    
-    * ``triton_int8kv``: 使用 int8 存储 kv cache，可增加 token 容量，使用 triton kernel
-    * ``ppl_int8kv``: 使用 int8 存储 kv cache，使用 ppl 快速 kernel
-    * ``ppl_fp16``: 使用 ppl 快速 fp16 解码注意力 kernel
-    * ``triton_flashdecoding``: 用于长上下文的 flashdecoding 模式，当前支持 llama llama2 qwen
-    * ``triton_gqa_attention``: 使用 GQA 的模型的快速 kernel
-    * ``triton_gqa_flashdecoding``: 使用 GQA 的模型的快速 flashdecoding kernel
-    * ``triton_fp8kv``: 使用 float8 存储 kv cache，目前仅用于 deepseek2
-    
-    需要阅读源代码以确认所有模型支持的具体模式
 
 调度参数
 --------
@@ -327,17 +311,9 @@ attention类型选择参数
 
     推理后端将为解码使用微批次重叠模式
     
-.. option:: --enable_flashinfer_prefill
+.. option:: --llm_kv_type
 
-    推理后端将为预填充使用 flashinfer 的注意力 kernel
-    
-.. option:: --enable_flashinfer_decode
-
-    推理后端将为解码使用 flashinfer 的注意力 kernel
-    
-.. option:: --enable_fa3
-
-    推理后端将为预填充和解码使用 fa3 注意力 kernel
+    推理后端使用什么类型的数据存储kv cache, 可选值为 "None", "int8kv", "int4kv", "fp8kv"
 
 .. option:: --disable_cudagraph
 
@@ -373,17 +349,14 @@ attention类型选择参数
 .. option:: --quant_type
 
     量化方法，可选值：
-    
-    * ``ppl-w4a16-128``
-    * ``flashllm-w6a16``
-    * ``ao-int4wo-[32,64,128,256]``
-    * ``ao-int8wo``
-    * ``ao-fp8w8a16``
-    * ``ao-fp6w6a16``
+
     * ``vllm-w8a8``
     * ``vllm-fp8w8a8``
     * ``vllm-fp8w8a8-b128``
+    * ``deepgemm-fp8w8a8-b128``
     * ``triton-fp8w8a8-block128``
+    * ``awq``
+    * ``awq_marlin``
     * ``none`` (默认)
 
 .. option:: --quant_cfg
@@ -395,13 +368,7 @@ attention类型选择参数
 .. option:: --vit_quant_type
 
     ViT 量化方法，可选值：
-    
-    * ``ppl-w4a16-128``
-    * ``flashllm-w6a16``
-    * ``ao-int4wo-[32,64,128,256]``
-    * ``ao-int8wo``
-    * ``ao-fp8w8a16``
-    * ``ao-fp6w6a16``
+
     * ``vllm-w8a8``
     * ``vllm-fp8w8a8``
     * ``none`` (默认)
diff --git a/docs/CN/source/tutorial/deepseek_deployment.rst b/docs/CN/source/tutorial/deepseek_deployment.rst
index 071d9405ab..720fe10e0b 100644
--- a/docs/CN/source/tutorial/deepseek_deployment.rst
+++ b/docs/CN/source/tutorial/deepseek_deployment.rst
@@ -33,12 +33,14 @@ LightLLM 支持以下几种部署模式：
     LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 8 \
-    --enable_fa3
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3
 
 **参数说明:**
 - `LOADWORKER=18`: 模型加载线程数，提高加载速度
 - `--tp 8`: 张量并行度，使用8个GPU
-- `--enable_fa3`: 启用 Flash Attention 3.0
+- `--llm_prefill_att_backend fa3`: 启用 Flash Attention 3.0
+- `--llm_decode_att_backend fa3`: 启用 Flash Attention 3.0
 - `--port 8088`: 服务端口
 
 1.2 单机 DP + EP 模式 (Data Parallel + Expert Parallel)
@@ -51,17 +53,20 @@ LightLLM 支持以下几种部署模式：
 .. code-block:: bash
 
     # H200 单机 DeepSeek-R1 DP + EP 模式
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
+    LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 8 \
     --dp 8 \
-    --enable_fa3
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
+    --enable_ep_moe
 
 **参数说明:**
-- `MOE_MODE=EP`: 设置专家并行模式
+- `--enable_ep_moe`: 设置专家并行模式
 - `--tp 8`: 张量并行度
 - `--dp 8`: 数据并行度，通常设置为与 tp 相同的值
-- `--enable_fa3`: 启用 Flash Attention 3.0
+- `--llm_prefill_att_backend fa3`: 启用 Flash Attention 3.0
+- `--llm_decode_att_backend fa3`: 启用 Flash Attention 3.0
 
 **可选优化参数:**
 - `--enable_prefill_microbatch_overlap`: 启用预填充微批次重叠
@@ -85,7 +90,8 @@ LightLLM 支持以下几种部署模式：
     LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 16 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --nnodes 2 \
     --node_rank 0 \
     --nccl_host $nccl_host \
@@ -101,7 +107,8 @@ LightLLM 支持以下几种部署模式：
     LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 16 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --nnodes 2 \
     --node_rank 1 \
     --nccl_host $nccl_host \
@@ -125,15 +132,16 @@ LightLLM 支持以下几种部署模式：
     # H200 多机 DeepSeek-R1 EP 模式 Node 0
     # 使用方法: sh multi_node_ep_node0.sh <nccl_host>
     export nccl_host=$1
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
+    LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 16 \
     --dp 16 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --nnodes 2 \
     --node_rank 0 \
     --nccl_host $nccl_host \
-    --nccl_port 2732
+    --nccl_port 2732 --enable_ep_moe
 
 **Node 1 启动命令:**
 
@@ -142,15 +150,16 @@ LightLLM 支持以下几种部署模式：
     # H200 多机 DeepSeek-R1 EP 模式 Node 1
     # 使用方法: sh multi_node_ep_node1.sh <nccl_host>
     export nccl_host=$1
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
+    LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 16 \
     --dp 16 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --nnodes 2 \
     --node_rank 1 \
     --nccl_host $nccl_host \
-    --nccl_port 2732
+    --nccl_port 2732 --enable_ep_moe
 
 **可选优化参数:**
 - `--enable_prefill_microbatch_overlap`: 启用预填充微批次重叠
@@ -187,7 +196,7 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
     export host=$1
     export pd_master_ip=$2
     nvidia-cuda-mps-control -d 
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+    LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "prefill" \
     --tp 8 \
@@ -195,10 +204,12 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
     --host $host \
     --port 8019 \
     --nccl_port 2732 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3  \
     --disable_cudagraph \
     --pd_master_ip $pd_master_ip \
-    --pd_master_port 60011
+    --pd_master_port 60011 \
+    --enable_ep_moe
     # 如果需要启用微批次重叠，可以取消注释以下行
     #--enable_prefill_microbatch_overlap
 
@@ -211,7 +222,7 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
     export host=$1
     export pd_master_ip=$2
     nvidia-cuda-mps-control -d
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+    LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "decode" \
     --tp 8 \
@@ -219,10 +230,12 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
     --host $host \
     --port 8121 \
     --nccl_port 12322 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --disable_cudagraph \
     --pd_master_ip $pd_master_ip \
-    --pd_master_port 60011
+    --pd_master_port 60011 \
+    --enable_ep_moe
     # 如果需要启用微批次重叠，可以取消注释以下行
     #--enable_decode_microbatch_overlap
 
@@ -279,7 +292,7 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
     export host=$1
     export config_server_host=$2
     nvidia-cuda-mps-control -d
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+    LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "prefill" \
     --host $host \
@@ -287,10 +300,12 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
     --tp 8 \
     --dp 8 \
     --nccl_port 2732 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --disable_cudagraph \
     --config_server_host $config_server_host \
-    --config_server_port 60088
+    --config_server_port 60088 \
+    --enable_ep_moe
     # 如果需要启用微批次重叠，可以取消注释以下行
     #--enable_prefill_microbatch_overlap
 
@@ -298,7 +313,7 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
     export host=$1
     export config_server_host=$2
     nvidia-cuda-mps-control -d
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+    LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "decode" \
     --host $host \
@@ -306,9 +321,11 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
     --nccl_port 12322 \
     --tp 8 \
     --dp 8 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --config_server_host $config_server_host \
-    --config_server_port 60088
+    --config_server_port 60088 \
+    --enable_ep_moe
     # 如果需要启用微批次重叠，可以取消注释以下行
     #--enable_decode_microbatch_overlap
 
diff --git a/docs/CN/source/tutorial/multi_level_cache_deployment.rst b/docs/CN/source/tutorial/multi_level_cache_deployment.rst
index 0446b07804..223b92dca3 100644
--- a/docs/CN/source/tutorial/multi_level_cache_deployment.rst
+++ b/docs/CN/source/tutorial/multi_level_cache_deployment.rst
@@ -66,7 +66,8 @@ LightLLM 的多级缓存系统采用分层设计:
         --model_dir /path/to/Qwen3-235B-A22B \
         --tp 8 \
         --graph_max_batch_size 500 \
-        --enable_fa3 \
+        --llm_prefill_att_backend fa3 \
+        --llm_decode_att_backend fa3  \
         --mem_fraction 0.88 \
         --enable_cpu_cache \
         --cpu_cache_storage_size 400 \
@@ -81,7 +82,7 @@ LightLLM 的多级缓存系统采用分层设计:
 - ``--model_dir``: 模型文件路径,支持本地路径或 HuggingFace 模型名称
 - ``--tp 8``: 张量并行度,使用 8 个 GPU 进行模型推理
 - ``--graph_max_batch_size 500``: CUDA Graph 最大批次大小,影响吞吐量和显存占用
-- ``--enable_fa3``: 启用 Flash Attention 3.0,提升注意力计算速度，也可以换成flashinfer后端性能更佳
+- ``--llm_prefill_att_backend fa3``: 启用 Flash Attention 3.0,提升注意力计算速度，也可以换成flashinfer后端性能更佳
 - ``--mem_fraction 0.88``: GPU 显存使用比例,建议设置为 0.88及以下
 
 CPU 缓存参数
@@ -130,7 +131,8 @@ CPU 缓存参数
         --model_dir /path/to/Qwen3-235B-A22B \
         --tp 8 \
         --graph_max_batch_size 500 \
-        --enable_fa3 \
+        --llm_prefill_att_backend fa3 \
+        --llm_decode_att_backend fa3  \
         --mem_fraction 0.88 \
         --enable_cpu_cache \
         --cpu_cache_storage_size 400 \
diff --git a/docs/CN/source/tutorial/reasoning_parser.rst b/docs/CN/source/tutorial/reasoning_parser.rst
index 547eb05d16..a9a0d09fe4 100644
--- a/docs/CN/source/tutorial/reasoning_parser.rst
+++ b/docs/CN/source/tutorial/reasoning_parser.rst
@@ -32,7 +32,8 @@ DeepSeek-R1
         --model_dir /path/to/DeepSeek-R1 \
         --reasoning_parser deepseek-r1 \
         --tp 8 \
-        --enable_fa3
+        --llm_prefill_att_backend fa3 \
+        --llm_decode_att_backend fa3
 
 DeepSeek-V3
 ~~~~~~~~~~~
diff --git a/docs/EN/source/getting_started/installation.rst b/docs/EN/source/getting_started/installation.rst
index 75fa714764..6439c48de3 100755
--- a/docs/EN/source/getting_started/installation.rst
+++ b/docs/EN/source/getting_started/installation.rst
@@ -24,16 +24,16 @@ The easiest way to install Lightllm is using the official image. You can directl
     $ docker pull ghcr.io/modeltc/lightllm:main
     $
     $ # Run，The current LightLLM service relies heavily on shared memory.
-    $ # Before starting, please make sure that you have allocated enough shared memory 
+    $ # Before starting, please make sure that you have allocated enough shared memory
     $ # in your Docker settings; otherwise, the service may fail to start properly.
     $ #
-    $ # 1. For text-only services, it is recommended to allocate more than 2GB of shared memory. 
+    $ # 1. For text-only services, it is recommended to allocate more than 2GB of shared memory.
     $ # If your system has sufficient RAM, allocating 16GB or more is recommended.
-    $ # 2.For multimodal services, it is recommended to allocate 16GB or more of shared memory. 
+    $ # 2.For multimodal services, it is recommended to allocate 16GB or more of shared memory.
     $ # You can adjust this value according to your specific requirements.
     $ #
-    $ # If you do not have enough shared memory available, you can try lowering 
-    $ # the --running_max_req_size parameter when starting the service. 
+    $ # If you do not have enough shared memory available, you can try lowering
+    $ # the --running_max_req_size parameter when starting the service.
     $ # This will reduce the number of concurrent requests, but also decrease shared memory usage.
     $ docker run -it --gpus all -p 8080:8080            \
     $   --shm-size 2g -v your_local_path:/data/         \
@@ -42,13 +42,13 @@ The easiest way to install Lightllm is using the official image. You can directl
 You can also manually build the image from source and run it:
 
 .. code-block:: console
-    
+
     $ # move into lightllm root dir
     $ cd /lightllm
     $ # Manually build the image
     $ docker build -t <image_name> -f ./docker/Dockerfile .
     $
-    $ # Run, 
+    $ # Run,
     $ docker run -it --gpus all -p 8080:8080            \
     $   --shm-size 2g -v your_local_path:/data/         \
     $   <image_name> /bin/bash
@@ -56,7 +56,7 @@ You can also manually build the image from source and run it:
 Or you can directly use the script to launch the image and run it with one click:
 
 .. code-block:: console
-    
+
     $ # View script parameters
     $ python tools/quick_launch_docker.py --help
 
@@ -84,6 +84,10 @@ You can also install Lightllm from source:
     $ # Install Lightllm dependencies (cuda 12.4)
     $ pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu124
     $
+    $ # Install Lightllm dependencies (Moore Threads GPU)
+    $ ./generate_requirements_musa.sh
+    $ pip install -r requirements-musa.txt
+    $
     $ # Install Lightllm
     $ python setup.py install
 
@@ -101,5 +105,5 @@ You can also install Lightllm from source:
     .. code-block:: console
 
         $ pip install -U --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/ triton-nightly --no-deps
-    
+
     For specific reasons, please refer to: `issue <https://github.com/triton-lang/triton/issues/3619>`_ and `fix PR <https://github.com/triton-lang/triton/pull/3638>`_
\ No newline at end of file
diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md
index 6127dffaf7..7417c39cfe 100755
--- a/docs/EN/source/models/add_new_model.md
+++ b/docs/EN/source/models/add_new_model.md
@@ -162,18 +162,6 @@ class BloomPreAndPostLayerWeight(PreAndPostLayerWeight):
                                                                  self.tp_rank_: split_vob_size * (self.tp_rank_ + 1), :])
             self.lm_head_weight_ = self.wte_weight_
         return
-    
-    def verify_load(self):
-        errors = "weights load not ok"
-        weights = [self.pre_norm_weight_, 
-                   self.pre_norm_bias_, 
-                   self.final_norm_weight_, 
-                   self.final_norm_bias_,
-                   self.wte_weight_,
-                   self.lm_head_weight_]
-        for i in range(len(weights)):
-            assert weights[i] is not None, "index:" + str(i) + " " + errors
-        return 
 
 ~~~
 
@@ -204,30 +192,6 @@ class BloomTransformerLayerWeight(TransformerLayerWeight):
         self._load_qkvo_weights(weights)
         self._load_ffn_weights(weights)
         return
-    
-    def verify_load(self):
-        errors = "weights load not ok"
-        weights = [self.att_norm_weight_,
-                   self.att_norm_bias_,
-                   self.q_weight_,
-                   self.k_weight_,
-                   self.v_weight_,
-                   self.q_bias_,
-                   self.k_bias_,
-                   self.v_bias_,
-                   self.o_weight_,
-                   self.o_bias_,
-
-                   self.ffn_norm_weight_,
-                   self.ffn_norm_bias_,
-                   self.ffn_1_weight_,
-                   self.ffn_1_bias_,
-                   self.ffn_2_weight_,
-                   self.ffn_2_bias_,
-                   ]
-        for i in range(len(weights)):
-            assert weights[i] is not None, "index:" + str(i) + " " + errors
-        return 
 
     def _load_qkvo_weights(self, weights):
         if f"h.{self.layer_num_}.input_layernorm.weight" in weights:
diff --git a/docs/EN/source/tutorial/api_server_args_zh.rst b/docs/EN/source/tutorial/api_server_args_zh.rst
index 1644bbab5f..42555ff9ec 100755
--- a/docs/EN/source/tutorial/api_server_args_zh.rst
+++ b/docs/EN/source/tutorial/api_server_args_zh.rst
@@ -183,23 +183,6 @@ Different Parallel Mode Setting Parameters
     When set to True, --nccl_host must equal config_server_host, --nccl_port must be unique for config_server,
     do not use the same nccl_port for different inference nodes, this will be a serious error
 
-Attention Type Selection Parameters
-------------------------------------
-
-.. option:: --mode
-
-    Model inference mode, can specify multiple values:
-    
-    * ``triton_int8kv``: Use int8 to store kv cache, can increase token capacity, uses triton kernel
-    * ``ppl_int8kv``: Use int8 to store kv cache, uses ppl fast kernel
-    * ``ppl_fp16``: Use ppl fast fp16 decode attention kernel
-    * ``triton_flashdecoding``: Flashdecoding mode for long context, currently supports llama llama2 qwen
-    * ``triton_gqa_attention``: Fast kernel for models using GQA
-    * ``triton_gqa_flashdecoding``: Fast flashdecoding kernel for models using GQA
-    * ``triton_fp8kv``: Use float8 to store kv cache, currently only used for deepseek2
-    
-    Need to read source code to confirm specific modes supported by all models 
-
 Scheduling Parameters
 ---------------------
 
@@ -325,18 +308,6 @@ Performance Optimization Parameters
 .. option:: --enable_decode_microbatch_overlap
 
     The inference backend will use microbatch overlap mode for decoding
-    
-.. option:: --enable_flashinfer_prefill
-
-    The inference backend will use flashinfer's attention kernel for prefill
-    
-.. option:: --enable_flashinfer_decode
-
-    The inference backend will use flashinfer's attention kernel for decoding
-    
-.. option:: --enable_fa3
-
-    The inference backend will use fa3 attention kernel for prefill and decoding
 
 .. option:: --disable_cudagraph
 
@@ -370,17 +341,14 @@ Quantization Parameters
 .. option:: --quant_type
 
     Quantization method, optional values:
-    
-    * ``ppl-w4a16-128``
-    * ``flashllm-w6a16``
-    * ``ao-int4wo-[32,64,128,256]``
-    * ``ao-int8wo``
-    * ``ao-fp8w8a16``
-    * ``ao-fp6w6a16``
+
     * ``vllm-w8a8``
     * ``vllm-fp8w8a8``
     * ``vllm-fp8w8a8-b128``
+    * ``deepgemm-fp8w8a8-b128``
     * ``triton-fp8w8a8-block128``
+    * ``awq``
+    * ``awq_marlin``
     * ``none`` (default)
 
 .. option:: --quant_cfg
@@ -392,13 +360,7 @@ Quantization Parameters
 .. option:: --vit_quant_type
 
     ViT quantization method, optional values:
-    
-    * ``ppl-w4a16-128``
-    * ``flashllm-w6a16``
-    * ``ao-int4wo-[32,64,128,256]``
-    * ``ao-int8wo``
-    * ``ao-fp8w8a16``
-    * ``ao-fp6w6a16``
+
     * ``vllm-w8a8``
     * ``vllm-fp8w8a8``
     * ``none`` (default)
diff --git a/docs/EN/source/tutorial/deepseek_deployment.rst b/docs/EN/source/tutorial/deepseek_deployment.rst
index 6098411be0..c908490afb 100755
--- a/docs/EN/source/tutorial/deepseek_deployment.rst
+++ b/docs/EN/source/tutorial/deepseek_deployment.rst
@@ -33,12 +33,13 @@ Suitable for deploying DeepSeek-R1 model on a single H200 node.
     LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 8 \
-    --enable_fa3
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3
 
 **Parameter Description:**
 - `LOADWORKER=18`: Model loading thread count, improves loading speed
 - `--tp 8`: Tensor parallelism, using 8 GPUs
-- `--enable_fa3`: Enable Flash Attention 3.0
+- `--llm_prefill_att_backend fa3`: Enable Flash Attention 3.0
 - `--port 8088`: Service port
 
 1.2 Single node DP + EP Mode (Data Parallel + Expert Parallel)
@@ -51,17 +52,18 @@ Suitable for expert parallelism deployment of MoE models like DeepSeek-V2/V3.
 .. code-block:: bash
 
     # H200 Single node DeepSeek-R1 DP + EP Mode
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
+    LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 8 \
     --dp 8 \
-    --enable_fa3
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
+    --enable_ep_moe
 
 **Parameter Description:**
-- `MOE_MODE=EP`: Set expert parallelism mode
+- `--enable_ep_moe`: Set expert parallelism mode
 - `--tp 8`: Tensor parallelism
 - `--dp 8`: Data parallelism, usually set to the same value as tp
-- `--enable_fa3`: Enable Flash Attention 3.0
 
 **Optional Optimization Parameters:**
 - `--enable_prefill_microbatch_overlap`: Enable prefill microbatch overlap
@@ -85,7 +87,8 @@ Suitable for deployment across multiple H200/H100 nodes.
     LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 16 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --nnodes 2 \
     --node_rank 0 \
     --nccl_host $nccl_host \
@@ -101,7 +104,8 @@ Suitable for deployment across multiple H200/H100 nodes.
     LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 16 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --nnodes 2 \
     --node_rank 1 \
     --nccl_host $nccl_host \
@@ -125,15 +129,16 @@ Suitable for deploying MoE models across multiple nodes.
     # H200 Multi-node DeepSeek-R1 EP Mode Node 0
     # Usage: sh multi_node_ep_node0.sh <nccl_host>
     export nccl_host=$1
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
+    LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 16 \
     --dp 16 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --nnodes 2 \
     --node_rank 0 \
     --nccl_host $nccl_host \
-    --nccl_port 2732
+    --nccl_port 2732 --enable_ep_moe
 
 **Node 1 Launch Command:**
 
@@ -142,15 +147,16 @@ Suitable for deploying MoE models across multiple nodes.
     # H200 Multi-node DeepSeek-R1 EP Mode Node 1
     # Usage: sh multi_node_ep_node1.sh <nccl_host>
     export nccl_host=$1
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
+    LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
     --model_dir /path/DeepSeek-R1 \
     --tp 16 \
     --dp 16 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --nnodes 2 \
     --node_rank 1 \
     --nccl_host $nccl_host \
-    --nccl_port 2732
+    --nccl_port 2732 --enable_ep_moe
 
 **Optional Optimization Parameters:**
 - `--enable_prefill_microbatch_overlap`: Enable prefill microbatch overlap
@@ -187,7 +193,7 @@ PD (Prefill-Decode) disaggregation mode separates prefill and decode stages for
     export host=$1
     export pd_master_ip=$2
     nvidia-cuda-mps-control -d 
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+    LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "prefill" \
     --tp 8 \
@@ -195,9 +201,11 @@ PD (Prefill-Decode) disaggregation mode separates prefill and decode stages for
     --host $host \
     --port 8019 \
     --nccl_port 2732 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --disable_cudagraph \
-    --pd_master_ip $pd_master_ip 
+    --pd_master_ip $pd_master_ip \
+    --enable_ep_moe
 
 **Step 3: Launch Decode Service**
 
@@ -208,7 +216,7 @@ PD (Prefill-Decode) disaggregation mode separates prefill and decode stages for
     export host=$1
     export pd_master_ip=$2
     nvidia-cuda-mps-control -d
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+    LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "decode" \
     --tp 8 \
@@ -216,10 +224,12 @@ PD (Prefill-Decode) disaggregation mode separates prefill and decode stages for
     --host $host \
     --port 8121 \
     --nccl_port 12322 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --disable_cudagraph \
     --pd_master_ip $pd_master_ip \
-    --pd_master_port 60011
+    --pd_master_port 60011 \
+    --enable_ep_moe
     # if you want to enable microbatch overlap, you can uncomment the following lines
     #--enable_decode_microbatch_overlap
 
@@ -276,7 +286,7 @@ Supports multiple PD Master nodes, providing better load balancing and high avai
     export host=$1
     export config_server_host=$2
     nvidia-cuda-mps-control -d
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+    LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "prefill" \
     --host $host \
@@ -284,10 +294,12 @@ Supports multiple PD Master nodes, providing better load balancing and high avai
     --tp 8 \
     --dp 8 \
     --nccl_port 2732 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --disable_cudagraph \
     --config_server_host $config_server_host \
-    --config_server_port 60088
+    --config_server_port 60088 \
+    --enable_ep_moe
     # if you want to enable microbatch overlap, you can uncomment the following lines
     #--enable_prefill_microbatch_overlap
 
@@ -295,7 +307,7 @@ Supports multiple PD Master nodes, providing better load balancing and high avai
     export host=$1
     export config_server_host=$2
     nvidia-cuda-mps-control -d
-    MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+    LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "decode" \
     --host $host \
@@ -303,9 +315,11 @@ Supports multiple PD Master nodes, providing better load balancing and high avai
     --nccl_port 12322 \
     --tp 8 \
     --dp 8 \
-    --enable_fa3 \
+    --llm_prefill_att_backend fa3 \
+    --llm_decode_att_backend fa3 \
     --config_server_host $config_server_host \
-    --config_server_port 60088
+    --config_server_port 60088 \
+    --enable_ep_moe
     # if you want to enable microbatch overlap, you can uncomment the following lines
     #--enable_decode_microbatch_overlap
 
diff --git a/docs/EN/source/tutorial/multi_level_cache_deployment.rst b/docs/EN/source/tutorial/multi_level_cache_deployment.rst
index bb8d943b87..6c99c351f0 100644
--- a/docs/EN/source/tutorial/multi_level_cache_deployment.rst
+++ b/docs/EN/source/tutorial/multi_level_cache_deployment.rst
@@ -66,7 +66,8 @@ Suitable for most scenarios, significantly increasing cache capacity while maint
         --model_dir /path/to/Qwen3-235B-A22B \
         --tp 8 \
         --graph_max_batch_size 500 \
-        --enable_fa3 \
+        --llm_prefill_att_backend fa3 \
+        --llm_decode_att_backend fa3 \
         --mem_fraction 0.88 \
         --enable_cpu_cache \
         --cpu_cache_storage_size 400 \
@@ -81,7 +82,7 @@ Basic Parameters
 - ``--model_dir``: Model file path, supports local path or HuggingFace model name
 - ``--tp 8``: Tensor parallelism degree, using 8 GPUs for model inference
 - ``--graph_max_batch_size 500``: CUDA Graph maximum batch size, affects throughput and memory usage
-- ``--enable_fa3``: Enable Flash Attention 3.0 to improve attention computation speed. You can also switch to flashinfer backend for better performance
+- ``--llm_prefill_att_backend fa3``: Enable Flash Attention 3.0 to improve attention computation speed. You can also switch to flashinfer backend for better performance
 - ``--mem_fraction 0.88``: GPU memory usage ratio, recommended to set to 0.88 or below
 
 CPU Cache Parameters
@@ -130,7 +131,8 @@ Suitable for ultra-long text or extremely high-concurrency scenarios, providing
         --model_dir /path/to/Qwen3-235B-A22B \
         --tp 8 \
         --graph_max_batch_size 500 \
-        --enable_fa3 \
+        --llm_prefill_att_backend fa3 \
+        --llm_decode_att_backend fa3 \
         --mem_fraction 0.88 \
         --enable_cpu_cache \
         --cpu_cache_storage_size 400 \
diff --git a/docs/EN/source/tutorial/reasoning_parser.rst b/docs/EN/source/tutorial/reasoning_parser.rst
index e76e093d63..56e61e6cd6 100644
--- a/docs/EN/source/tutorial/reasoning_parser.rst
+++ b/docs/EN/source/tutorial/reasoning_parser.rst
@@ -32,7 +32,8 @@ DeepSeek-R1
         --model_dir /path/to/DeepSeek-R1 \
         --reasoning_parser deepseek-r1 \
         --tp 8 \
-        --enable_fa3
+        --llm_prefill_att_backend fa3 \
+        --llm_decode_att_backend fa3
 
 DeepSeek-V3
 ~~~~~~~~~~~
diff --git a/generate_requirements_musa.sh b/generate_requirements_musa.sh
new file mode 100755
index 0000000000..f5bfb8ff83
--- /dev/null
+++ b/generate_requirements_musa.sh
@@ -0,0 +1,105 @@
+#!/bin/bash
+# Script to generate requirements-musa.txt from requirements.txt
+# MUSA is not compatible with CUDA packages, so they need to be removed
+# Torch-related packages are pre-installed in the MUSA docker container
+
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+INPUT_FILE="${SCRIPT_DIR}/requirements.txt"
+OUTPUT_FILE="${SCRIPT_DIR}/requirements-musa.txt"
+
+if [ ! -f "$INPUT_FILE" ]; then
+    echo "Error: requirements.txt not found at $INPUT_FILE"
+    exit 1
+fi
+
+echo "Generating requirements-musa.txt from requirements.txt..."
+
+# Define patterns to remove (CUDA-specific packages)
+# These packages are not compatible with MUSA
+CUDA_PACKAGES=(
+    "^cupy"           # cupy-cuda12x and similar
+    "^cuda_bindings"  # CUDA bindings
+    "^nixl"           # NIXL (NVIDIA Inter-node eXchange Library)
+    "^flashinfer"     # flashinfer-python (CUDA-specific attention kernel)
+    "^sgl-kernel"     # SGL kernel (CUDA-specific)
+)
+
+# Define torch-related packages (pre-installed in MUSA container, remove version pins)
+TORCH_PACKAGES=(
+    "^torch=="
+    "^torchvision=="
+)
+
+# Create the output file with a header comment
+cat > "$OUTPUT_FILE" << 'EOF'
+# Requirements for MUSA (Moore Threads GPU)
+# Auto-generated from requirements.txt by generate_requirements_musa.sh
+# CUDA-specific packages have been removed
+# Torch-related packages have version pins removed (pre-installed in MUSA container)
+
+EOF
+
+# Process the requirements file
+while IFS= read -r line || [ -n "$line" ]; do
+    # Skip empty lines and comments (but keep them in output)
+    if [[ -z "$line" || "$line" =~ ^[[:space:]]*# ]]; then
+        echo "$line" >> "$OUTPUT_FILE"
+        continue
+    fi
+
+    # Extract package name (before ==, >=, <=, ~=, etc.)
+    pkg_name=$(echo "$line" | sed -E 's/^([a-zA-Z0-9_-]+).*/\1/')
+
+    # Check if this is a CUDA package to skip
+    skip=false
+    for pattern in "${CUDA_PACKAGES[@]}"; do
+        if [[ "$pkg_name" =~ $pattern ]]; then
+            echo "  Removing CUDA package: $line"
+            skip=true
+            break
+        fi
+    done
+
+    if $skip; then
+        continue
+    fi
+
+    # Check if this is a torch-related package (remove version pin)
+    for pattern in "${TORCH_PACKAGES[@]}"; do
+        if [[ "$line" =~ $pattern ]]; then
+            # Remove version pin, keep just the package name
+            pkg_only=$(echo "$line" | sed -E 's/==.*//')
+            echo "  Unpinning version for: $pkg_only (pre-installed in MUSA container)"
+            echo "$pkg_only" >> "$OUTPUT_FILE"
+            skip=true
+            break
+        fi
+    done
+
+    if $skip; then
+        continue
+    fi
+
+    # Keep the package as-is
+    echo "$line" >> "$OUTPUT_FILE"
+
+done < "$INPUT_FILE"
+
+# Add MUSA-specific packages at the end
+cat >> "$OUTPUT_FILE" << 'EOF'
+
+# MUSA-specific packages
+torch_musa
+torchada
+EOF
+
+echo ""
+echo "Successfully generated: $OUTPUT_FILE"
+echo ""
+echo "Summary of changes:"
+echo "  - Removed CUDA-specific packages: cupy-cuda12x, cuda_bindings, nixl, flashinfer-python, sgl-kernel"
+echo "  - Unpinned torch-related packages: torch, torchvision (pre-installed in MUSA container)"
+echo "  - Added MUSA-specific packages: torch_musa, torchada"
+
diff --git a/lightllm/__init__.py b/lightllm/__init__.py
index e69de29bb2..e9ba6f3041 100644
--- a/lightllm/__init__.py
+++ b/lightllm/__init__.py
@@ -0,0 +1,4 @@
+from lightllm.utils.device_utils import is_musa
+
+if is_musa():
+    import torchada  # noqa: F401
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..67dd3852c5
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..555386ebdd
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..6f92439c1c
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..67dd3852c5
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..555386ebdd
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..6f92439c1c
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..7f69e86a86
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}, "32": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 3}, "128": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..7d8dc868c3
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 5}}, "8192": {"8": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 10}, "32": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 5}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..3e543b2ea3
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..7f69e86a86
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}, "32": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 3}, "128": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..7d8dc868c3
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 5}}, "8192": {"8": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 10}, "32": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 5}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..3e543b2ea3
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=16,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..a3b0edde6d
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..22d1ce6f69
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 1}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 2}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..328fcec837
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 16, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 1}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..a3b0edde6d
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..22d1ce6f69
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 1}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 2}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..328fcec837
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 16, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 1}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..4c4ae86241
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..61884a9375
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 1}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..037bfd2913
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 1}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..4c4ae86241
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..61884a9375
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 1}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..037bfd2913
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=2,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 1}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..e2028e2d2a
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..7e99dc1be2
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 1}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 1}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..5d6b46dda8
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 4}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..e2028e2d2a
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 16, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..7e99dc1be2
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 1}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 1}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..5d6b46dda8
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 4}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..7795b47e72
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 4}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..ff4d6efd49
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 4}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..1d8ca6967b
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..7795b47e72
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 4}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..ff4d6efd49
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 4}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..1d8ca6967b
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=4,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..2f1cd5dfd5
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..a369088bfb
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..cb4e6a0d3e
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..2f1cd5dfd5
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..a369088bfb
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..cb4e6a0d3e
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..60827b791e
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 4}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..7b42cad466
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..9bb49d70b7
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..60827b791e
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 4}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..7b42cad466
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..9bb49d70b7
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=5,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..bd3d1c418b
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..f1b3539f5d
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 1}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..e12c05b966
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..bd3d1c418b
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..f1b3539f5d
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 1}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}, "256": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 2, "num_stages": 2}, "128": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..e12c05b966
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=4,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}, "128": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 4}, "256": {"BLOCK_N": 64, "num_warps": 2, "num_stages": 3}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..c83dca52d2
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..59a3e1051c
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..5b7c4eaa9f
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..c83dca52d2
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..59a3e1051c
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}, "8192": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "128": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "256": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
new file mode 100644
index 0000000000..5b7c4eaa9f
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage1:v2/{block_seq=256,gqa_group_size=8,max_batch_group_size=8,out_dtype=torch.float16,q_head_dim=128}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}, "8192": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "128": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}, "256": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 5}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..abd760af04
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1 @@
+{"32": {"8": {"BLOCK_N": 32, "num_warps": 16, "num_stages": 3}, "16": {"BLOCK_N": 16, "num_warps": 8, "num_stages": 5}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 2}, "64": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 2}}, "64": {"8": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 2}, "16": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "64": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 2}}, "128": {"8": {"BLOCK_N": 32, "num_warps": 16, "num_stages": 2}, "16": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 2}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 4}, "64": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}, "256": {"8": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 1}, "16": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 9}, "32": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "64": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 2}}}
\ No newline at end of file
diff --git a/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..a560ce9e1a
--- /dev/null
+++ b/lightllm/common/all_kernel_configs/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1 @@
+{"32": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}, "16": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 2}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 11}, "64": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}, "64": {"8": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 7}, "16": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 10}, "32": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}, "64": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 4}}, "128": {"8": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 3}, "16": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 7}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "64": {"BLOCK_N": 16, "num_warps": 2, "num_stages": 3}}, "256": {"8": {"BLOCK_N": 64, "num_warps": 16, "num_stages": 2}, "16": {"BLOCK_N": 64, "num_warps": 8, "num_stages": 5}, "32": {"BLOCK_N": 64, "num_warps": 4, "num_stages": 3}, "64": {"BLOCK_N": 32, "num_warps": 4, "num_stages": 2}}}
\ No newline at end of file
diff --git a/lightllm/common/basemodel/attention/__init__.py b/lightllm/common/basemodel/attention/__init__.py
new file mode 100644
index 0000000000..80df545498
--- /dev/null
+++ b/lightllm/common/basemodel/attention/__init__.py
@@ -0,0 +1,18 @@
+from .base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
+from .triton.fp import TritonAttBackend
+from .triton.int4kv import Int4kvTritonAttBackend
+from .triton.int8kv import Int8kvTritonAttBackend
+from .triton.mla import MlaTritonAttBackend
+from .fa3.fp import Fa3AttBackend
+from .fa3.fp8 import Fp8Fa3AttBackend
+from .fa3.mla import MlaFa3AttBackend
+from .flashinfer.fp8 import Fp8FlashInferAttBackend
+from .flashinfer.fp import FlashInferAttBackend
+from .flashinfer.mla import MlaFlashInferAttBackend
+
+from .create_utils import (
+    get_prefill_att_backend_class,
+    get_decode_att_backend_class,
+    get_mla_prefill_att_backend_class,
+    get_mla_decode_att_backend_class,
+)
diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py
new file mode 100644
index 0000000000..859d97ca84
--- /dev/null
+++ b/lightllm/common/basemodel/attention/base_att.py
@@ -0,0 +1,117 @@
+import torch
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional, TYPE_CHECKING, Tuple, Union, Dict
+
+if TYPE_CHECKING:
+    from lightllm.common.basemodel.basemodel import TpPartBaseModel
+    from lightllm.common.basemodel.infer_struct import InferStateInfo
+
+
+class BaseAttBackend:
+    """
+    用于创建支持各种不同的AttBackend, 如 fa3, flashinfer, triton 实现等，
+    这个是单列模式, 每种backend只有一个实例
+    """
+
+    _instances = {}
+
+    def __new__(cls, *args, **kwargs):
+        """
+        重写__new__方法实现单例模式
+        """
+        # 检查是否已经有该类的实例
+        if cls not in cls._instances:
+            # 创建新实例并存储
+            instance = super().__new__(cls)
+            cls._instances[cls] = instance
+        # 返回已有的实例
+        return cls._instances[cls]
+
+    def __init__(self, model: "TpPartBaseModel"):
+        self.model = model
+
+    def create_att_prefill_state(self) -> "BasePrefillAttState":
+        raise NotImplementedError("not impl")
+
+    def create_att_decode_state(self) -> "BaseDecodeAttState":
+        raise NotImplementedError("not impl")
+
+    def _find_layer_index(
+        self, k: torch.Tensor, v: torch.Tensor, att_state: Union["BasePrefillAttState", "BaseDecodeAttState"]
+    ) -> int:
+        kv_buffer = att_state.infer_state.mem_manager.kv_buffer
+        layer_count = len(kv_buffer)
+        find_dict = {kv_buffer[i].data_ptr(): i for i in range(layer_count)}
+        key = min(k.data_ptr(), v.data_ptr())
+        assert key in find_dict
+        return find_dict[key]
+
+
+@dataclass
+class AttControl:
+    """
+    prefill_att 和 decode_att 的入参，用于控制att backend 内部的行为, 选择正确的att 实现。
+    """
+
+    use_alibi: bool = False
+    tp_alibi: torch.Tensor = None
+    use_sliding_window: bool = False
+    sliding_window: Tuple[int, int] = (-1, -1)
+    use_att_sink: bool = False
+    sink_weight: torch.Tensor = None
+    # mla 专用传参项
+    mla_prefill: bool = False
+    mla_prefill_dict: Dict = None
+    mla_decode: bool = False
+    mla_decode_dict: Dict = None
+
+
+@dataclass
+class BasePrefillAttState(ABC):
+
+    backend: BaseAttBackend = None
+    infer_state: "InferStateInfo" = None
+
+    @abstractmethod
+    def init_state(self):
+        pass
+
+    @abstractmethod
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        raise NotImplementedError("not impl")
+
+
+@dataclass
+class BaseDecodeAttState(ABC):
+    backend: BaseAttBackend = None
+    infer_state: "InferStateInfo" = None
+
+    @abstractmethod
+    def init_state(self):
+        pass
+
+    def copy_for_decode_cuda_graph(self, new_state: "BaseDecodeAttState"):
+        for attr_name, attr_value in vars(new_state).items():
+            if isinstance(attr_value, torch.Tensor):
+                attr_ = getattr(self, attr_name, None)
+                if attr_ is not None and attr_.data_ptr() != attr_value.data_ptr():
+                    attr_.copy_(attr_value, non_blocking=True)
+
+    @abstractmethod
+    def decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        pass
diff --git a/lightllm/common/basemodel/attention/create_utils.py b/lightllm/common/basemodel/attention/create_utils.py
new file mode 100644
index 0000000000..39e32ac635
--- /dev/null
+++ b/lightllm/common/basemodel/attention/create_utils.py
@@ -0,0 +1,80 @@
+from lightllm.utils.envs_utils import get_env_start_args
+from .base_att import BaseAttBackend
+from .triton.fp import TritonAttBackend
+from .triton.int4kv import Int4kvTritonAttBackend
+from .triton.int8kv import Int8kvTritonAttBackend
+from .triton.mla import MlaTritonAttBackend
+from .fa3.fp import Fa3AttBackend
+from .fa3.fp8 import Fp8Fa3AttBackend
+from .fa3.mla import MlaFa3AttBackend
+from .flashinfer.fp8 import Fp8FlashInferAttBackend
+from .flashinfer.fp import FlashInferAttBackend
+from .flashinfer.mla import MlaFlashInferAttBackend
+
+data_type_to_backend = {
+    "None": {
+        "triton": TritonAttBackend,
+        "fa3": Fa3AttBackend,
+        "flashinfer": FlashInferAttBackend,
+    },
+    "int4kv": {
+        "triton": Int4kvTritonAttBackend,
+        "fa3": Fp8Fa3AttBackend,
+        "flashinfer": Fp8FlashInferAttBackend,
+    },
+    "int8kv": {
+        "triton": Int8kvTritonAttBackend,
+        "fa3": Fp8Fa3AttBackend,
+        "flashinfer": Fp8FlashInferAttBackend,
+    },
+}
+
+mla_data_type_to_backend = {
+    "None": {
+        "triton": MlaTritonAttBackend,
+        "fa3": MlaFa3AttBackend,
+        "flashinfer": MlaFlashInferAttBackend,
+    },
+}
+
+
+def get_prefill_att_backend_class(index=0) -> BaseAttBackend:
+    args = get_env_start_args()
+    llm_dtype = args.llm_kv_type
+    backend_str = args.llm_prefill_att_backend[index]
+    if backend_str != "None":
+        return data_type_to_backend[llm_dtype][backend_str]
+    else:
+        # 根据环境自动选择最好的
+        raise NotImplementedError(f"error")
+
+
+def get_decode_att_backend_class(index=0) -> BaseAttBackend:
+    args = get_env_start_args()
+    llm_dtype = args.llm_kv_type
+    backend_str = args.llm_decode_att_backend[index]
+    if backend_str != "None":
+        return data_type_to_backend[llm_dtype][backend_str]
+    else:
+        # 根据环境自动选择最好的
+        raise NotImplementedError(f"error")
+
+
+def get_mla_prefill_att_backend_class(index=0) -> BaseAttBackend:
+    args = get_env_start_args()
+    llm_dtype = args.llm_kv_type
+    backend_str = args.llm_prefill_att_backend[index]
+    if backend_str != "None":
+        return mla_data_type_to_backend[llm_dtype][backend_str]
+    else:
+        raise NotImplementedError(f"error")
+
+
+def get_mla_decode_att_backend_class(index=0) -> BaseAttBackend:
+    args = get_env_start_args()
+    llm_dtype = args.llm_kv_type
+    backend_str = args.llm_decode_att_backend[index]
+    if backend_str != "None":
+        return mla_data_type_to_backend[llm_dtype][backend_str]
+    else:
+        raise NotImplementedError(f"error")
diff --git a/lightllm/common/fused_moe/__init__.py b/lightllm/common/basemodel/attention/fa3/__init__.py
similarity index 100%
rename from lightllm/common/fused_moe/__init__.py
rename to lightllm/common/basemodel/attention/fa3/__init__.py
diff --git a/lightllm/common/basemodel/attention/fa3/fp.py b/lightllm/common/basemodel/attention/fa3/fp.py
new file mode 100644
index 0000000000..952bb39d91
--- /dev/null
+++ b/lightllm/common/basemodel/attention/fa3/fp.py
@@ -0,0 +1,243 @@
+import dataclasses
+import torch
+from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
+from typing import Optional, TYPE_CHECKING
+from lightllm.utils.dist_utils import get_current_device_id
+from lightllm.utils.sgl_utils import flash_attn_with_kvcache
+from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.common.basemodel.triton_kernel.fa3_utils import page_table_copy
+from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor
+
+
+class Fa3AttBackend(BaseAttBackend):
+    def __init__(self, model):
+        super().__init__(model=model)
+        self.get_page_table_buffer()  # init
+
+    def get_page_table_buffer(self):
+        """
+        用于减少 decode graph 捕获的时候, 造成显存二次方增长的情况.
+        """
+        model = self.model
+        if not hasattr(self, "_shared_page_table_buffer"):
+            self._shared_page_table_buffer = [
+                torch.empty(model.graph_max_batch_size * model.graph_max_len_in_batch, dtype=torch.int32).to(
+                    get_current_device_id()
+                ),
+                torch.empty(model.graph_max_batch_size * model.graph_max_len_in_batch, dtype=torch.int32).to(
+                    get_current_device_id()
+                ),
+            ]
+        return self._shared_page_table_buffer
+
+    def create_att_prefill_state(self, infer_state) -> "Fa3PrefillAttState":
+        return Fa3PrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state) -> "Fa3DecodeAttState":
+        return Fa3DecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class Fa3PrefillAttState(BasePrefillAttState):
+    cu_seqlens_q: torch.Tensor = None
+    cu_seqlens_k: torch.Tensor = None
+    page_table: torch.Tensor = None
+
+    def init_state(self):
+        self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int()
+        self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int()
+        self.page_table = torch.empty(
+            (self.infer_state.batch_size, self.infer_state.max_kv_seq_len),
+            dtype=torch.int32,
+            device=self.infer_state.input_ids.device,
+        )
+        self.page_table.copy_(
+            self.infer_state.req_manager.req_to_token_indexs[
+                self.infer_state.b_req_idx, : self.infer_state.max_kv_seq_len
+            ]
+        )
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert att_control.use_alibi is False
+        return self._nomarl_prefill_att(
+            q=q,
+            k=k,
+            v=v,
+            att_control=att_control,
+            alloc_func=alloc_func,
+        )
+
+    def _nomarl_prefill_att(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, att_control: AttControl, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        self.backend: Fa3AttBackend = self.backend  # for typing
+
+        if att_control.use_sliding_window:
+            window_size = att_control.sliding_window
+        else:
+            window_size = (-1, -1)
+
+        if att_control.use_att_sink:
+            sink_weight: torch.Tensor = att_control.sink_weight
+        else:
+            sink_weight = None
+
+        k_descale, v_descale = None, None  # disable quantization
+        Lq = q.shape[-1]
+        sm_scale = 1.0 / (Lq ** 0.5)
+        o = flash_attn_with_kvcache(
+            q=q,
+            k_cache=k.view(k.shape[0], 1, k.shape[1], k.shape[2]),
+            v_cache=v.view(v.shape[0], 1, v.shape[1], v.shape[2]),
+            page_table=self.page_table,
+            cache_seqlens=self.infer_state.b_seq_len,
+            cu_seqlens_q=self.cu_seqlens_q,
+            cu_seqlens_k_new=self.cu_seqlens_k,
+            max_seqlen_q=self.infer_state.max_q_seq_len,
+            softmax_scale=sm_scale,
+            causal=True,
+            window_size=window_size,
+            softcap=0.0,
+            k_descale=k_descale,
+            v_descale=v_descale,
+            return_softmax_lse=False,
+            sinks=sink_weight,
+        )
+        return o
+
+
+@dataclasses.dataclass
+class Fa3DecodeAttState(BaseDecodeAttState):
+    cu_seqlens_q: torch.Tensor = None
+    cu_seqlens_k: torch.Tensor = None
+    page_table: torch.Tensor = None
+    b_att_seq_len: torch.Tensor = None
+    # 在是否开启mtp 的不同模式下，其设置不同的值，可以加速算子的运行。
+    decode_max_q_seq_len: int = None
+
+    def init_state(self):
+        self.backend: Fa3AttBackend = self.backend
+        args_mtp_step = get_env_start_args().mtp_step
+
+        if args_mtp_step > 0:
+            # 修正 mtp 在 fa3 下的输入。
+            mtp_size = args_mtp_step + 1
+            b_q_seq_len = torch.full(
+                (self.infer_state.b_seq_len.shape[0] // mtp_size,),
+                fill_value=mtp_size,
+                dtype=torch.int32,
+                device=self.infer_state.b_seq_len.device,
+            )
+            b_kv_seq_len = self.infer_state.b_seq_len[mtp_size - 1 :: mtp_size]
+            b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len)
+            self.cu_seqlens_q = b1_cu_q_seq_len.int()
+            self.cu_seqlens_k = b1_cu_kv_seq_len.int()
+        else:
+            self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int()
+            self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int()
+
+        att_batch_size = self.infer_state.batch_size // (args_mtp_step + 1)
+        assert self.infer_state.batch_size % (args_mtp_step + 1) == 0
+
+        model = self.backend.model
+        # 可以使用 cuda graph的时候从 buffer中申请
+        if (
+            self.infer_state.batch_size <= model.graph_max_batch_size
+            and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch
+        ):
+            page_buffer = self.backend.get_page_table_buffer()
+            self.page_table = page_buffer[self.infer_state.microbatch_index][
+                : att_batch_size * model.graph_max_len_in_batch
+            ].reshape(att_batch_size, model.graph_max_len_in_batch)
+        else:
+            self.page_table = torch.empty(
+                (att_batch_size, self.infer_state.max_kv_seq_len),
+                dtype=torch.int32,
+                device=self.infer_state.input_ids.device,
+            )
+
+        if args_mtp_step > 0:
+            page_table_copy(
+                page_table=self.page_table[:, : self.infer_state.max_kv_seq_len],
+                req_to_token_indexs=model.req_manager.req_to_token_indexs,
+                b_req_idx=self.infer_state.b_req_idx[args_mtp_step :: (args_mtp_step + 1)],
+            )
+            self.b_att_seq_len = self.infer_state.b_seq_len[args_mtp_step :: (args_mtp_step + 1)].contiguous()
+            self.decode_max_q_seq_len = args_mtp_step + 1
+        else:
+            page_table_copy(
+                page_table=self.page_table[:, : self.infer_state.max_kv_seq_len],
+                req_to_token_indexs=model.req_manager.req_to_token_indexs,
+                b_req_idx=self.infer_state.b_req_idx,
+            )
+            self.b_att_seq_len = self.infer_state.b_seq_len
+            self.decode_max_q_seq_len = 1
+        return
+
+    def copy_for_decode_cuda_graph(self, new_state: "Fa3DecodeAttState"):
+        super().copy_for_decode_cuda_graph(new_state)
+
+    def decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ):
+        assert att_control.use_alibi is False
+        return self._normal_decode_att(
+            q=q,
+            k=k,
+            v=v,
+            att_control=att_control,
+            alloc_func=alloc_func,
+        )
+
+    def _normal_decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl,
+        alloc_func=torch.empty,
+    ):
+        if att_control.use_sliding_window:
+            window_size = att_control.sliding_window
+        else:
+            window_size = (-1, -1)
+
+        if att_control.use_att_sink:
+            sink_weight: torch.Tensor = att_control.sink_weight
+        else:
+            sink_weight = None
+
+        k_descale, v_descale = None, None  # disable quantization
+        Lq = q.shape[-1]
+        sm_scale = 1.0 / (Lq ** 0.5)
+        o = flash_attn_with_kvcache(
+            q=q,
+            k_cache=k.view(k.shape[0], 1, k.shape[1], k.shape[2]),
+            v_cache=v.view(v.shape[0], 1, v.shape[1], v.shape[2]),
+            page_table=self.page_table,
+            cache_seqlens=self.b_att_seq_len,
+            cu_seqlens_q=self.cu_seqlens_q,
+            cu_seqlens_k_new=self.cu_seqlens_k,
+            max_seqlen_q=self.decode_max_q_seq_len,
+            softmax_scale=sm_scale,
+            causal=True,
+            window_size=window_size,
+            softcap=0.0,
+            k_descale=k_descale,
+            v_descale=v_descale,
+            return_softmax_lse=False,
+            sinks=sink_weight,
+        )
+        return o
diff --git a/lightllm/common/basemodel/attention/fa3/fp8.py b/lightllm/common/basemodel/attention/fa3/fp8.py
new file mode 100644
index 0000000000..12b2b0dfa8
--- /dev/null
+++ b/lightllm/common/basemodel/attention/fa3/fp8.py
@@ -0,0 +1,221 @@
+import dataclasses
+import torch
+from ..base_att import AttControl
+from typing import Optional, TYPE_CHECKING
+from lightllm.utils.sgl_utils import flash_attn_with_kvcache
+from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant import q_per_head_fp8_quant
+from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops
+from typing import Union
+from .fp import Fa3AttBackend, Fa3PrefillAttState, Fa3DecodeAttState
+
+if HAS_VLLM:
+    scaled_fp8_quant = vllm_ops.scaled_fp8_quant
+else:
+    scaled_fp8_quant = None
+
+
+class Fp8Fa3AttBackend(Fa3AttBackend):
+    def __init__(self, model):
+        super().__init__(model=model)
+
+    def create_att_prefill_state(self, infer_state) -> "Fp8Fa3PrefillAttState":
+        return Fp8Fa3PrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state) -> "Fp8Fa3DecodeAttState":
+        return Fp8Fa3DecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class Fp8Fa3PrefillAttState(Fa3PrefillAttState):
+    # 临时共享变量
+    mid_token_batch_ids: torch.Tensor = None
+    k_descale: torch.Tensor = None
+    v_descale: torch.Tensor = None
+
+    def init_state(self):
+        super().init_state()
+        device = self.infer_state.input_ids.device
+        batch_size = self.infer_state.batch_size
+        mem_manager = self.backend.model.mem_manager
+
+        offline_scales: torch.Tensor = mem_manager.scales
+        head_num = mem_manager.head_num
+        self.mid_token_batch_ids = torch.repeat_interleave(
+            torch.arange(batch_size, device=device), self.infer_state.b_q_seq_len
+        )
+        # 为了减少推理计算量，在推理外部初始化k_descale和v_descale
+        self.k_descale = (
+            offline_scales[:, :head_num].view(-1, 1, head_num).expand(offline_scales.shape[0], batch_size, head_num)
+            if offline_scales is not None
+            else torch.ones(
+                (mem_manager.layer_num, batch_size, head_num),
+                dtype=torch.float32,
+                device=device,
+            )
+        )
+        self.v_descale = (
+            offline_scales[:, head_num:].view(-1, 1, head_num).expand(offline_scales.shape[0], batch_size, head_num)
+            if offline_scales is not None
+            else torch.ones(
+                (mem_manager.layer_num, batch_size, head_num),
+                dtype=torch.float32,
+                device=device,
+            )
+        )
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        return self._fp8_prefill_att(
+            q=q,
+            k=k,
+            v=v,
+            alloc_func=alloc_func,
+        )
+
+    def _fp8_prefill_att(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        self.backend: Fp8Fa3AttBackend = self.backend  # for typing
+
+        q, q_scale = q_per_head_fp8_quant(
+            q,
+            self.infer_state.b_seq_len,
+            self.cu_seqlens_q,
+            self.mid_token_batch_ids,
+        )
+        k_head_num = k.shape[1]
+        k_head_dim = k.shape[2]
+        cache_k = k.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn)
+        cache_v = v.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn)
+        layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, att_state=self)
+        o = flash_attn_with_kvcache(
+            q=q,
+            k_cache=cache_k,
+            v_cache=cache_v,
+            page_table=self.page_table,
+            cache_seqlens=self.infer_state.b_seq_len,
+            cu_seqlens_q=self.cu_seqlens_q,
+            cu_seqlens_k_new=self.cu_seqlens_k,
+            max_seqlen_q=self.infer_state.max_q_seq_len,
+            causal=True,
+            window_size=(-1, -1),
+            softcap=0.0,
+            q_descale=q_scale,
+            k_descale=self.k_descale[layer_index],
+            v_descale=self.v_descale[layer_index],
+            return_softmax_lse=False,
+        )
+        return o
+
+
+@dataclasses.dataclass
+class Fp8Fa3DecodeAttState(Fa3DecodeAttState):
+    k_descale: torch.Tensor = None
+    v_descale: torch.Tensor = None
+
+    def init_state(self):
+        super().init_state()
+        self.backend: Fp8Fa3AttBackend = self.backend
+
+        args_mtp_step = get_env_start_args().mtp_step
+        att_batch_size = self.infer_state.batch_size // (args_mtp_step + 1)
+        assert self.infer_state.batch_size % (args_mtp_step + 1) == 0
+
+        device = self.infer_state.input_ids.device
+        batch_size = att_batch_size
+        mem_manager = self.backend.model.mem_manager
+
+        offline_scales: torch.Tensor = mem_manager.scales
+        head_num = mem_manager.head_num
+
+        # 为了减少推理计算量，在推理外部初始化k_descale和v_descale
+        self.k_descale = (
+            offline_scales[:, :head_num].view(-1, 1, head_num).expand(offline_scales.shape[0], batch_size, head_num)
+            if offline_scales is not None
+            else torch.ones(
+                (mem_manager.layer_num, batch_size, head_num),
+                dtype=torch.float32,
+                device=device,
+            )
+        )
+        self.v_descale = (
+            offline_scales[:, head_num:].view(-1, 1, head_num).expand(offline_scales.shape[0], batch_size, head_num)
+            if offline_scales is not None
+            else torch.ones(
+                (mem_manager.layer_num, batch_size, head_num),
+                dtype=torch.float32,
+                device=device,
+            )
+        )
+        return
+
+    def copy_for_decode_cuda_graph(self, new_state: "Fp8Fa3DecodeAttState"):
+        super().copy_for_decode_cuda_graph(new_state)
+
+    def decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ):
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        return self._fp8_decode_att(
+            q=q,
+            k=k,
+            v=v,
+            alloc_func=alloc_func,
+        )
+
+    def _fp8_decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        alloc_func=torch.empty,
+    ):
+        k_head_num = k.shape[1]
+        k_head_dim = k.shape[2]
+
+        cache_k = k.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn)
+        cache_v = v.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn)
+
+        layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, att_state=self)
+
+        q_head_num = q.shape[1]
+        q, q_scale = scaled_fp8_quant(q.view(q.shape[0] * k_head_num, -1), use_per_token_if_dynamic=True)
+        o = flash_attn_with_kvcache(
+            q=q.view(-1, q_head_num, k_head_dim),
+            k_cache=cache_k,
+            v_cache=cache_v,
+            page_table=self.page_table,
+            cache_seqlens=self.infer_state.b_seq_len,
+            cu_seqlens_q=self.cu_seqlens_q,
+            cu_seqlens_k_new=self.cu_seqlens_k,
+            max_seqlen_q=self.decode_max_q_seq_len,
+            causal=False,
+            window_size=(-1, -1),
+            softcap=0.0,
+            q_descale=q_scale.view(self.infer_state.batch_size, k_head_num),
+            k_descale=self.k_descale[layer_index],
+            v_descale=self.v_descale[layer_index],
+            return_softmax_lse=False,
+        )
+        return o
diff --git a/lightllm/common/basemodel/attention/fa3/mla.py b/lightllm/common/basemodel/attention/fa3/mla.py
new file mode 100644
index 0000000000..9a10457b12
--- /dev/null
+++ b/lightllm/common/basemodel/attention/fa3/mla.py
@@ -0,0 +1,229 @@
+import dataclasses
+import torch
+from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
+from typing import Optional, TYPE_CHECKING, Tuple
+from lightllm.utils.dist_utils import get_current_device_id
+from lightllm.utils.sgl_utils import flash_attn_with_kvcache
+from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.common.basemodel.triton_kernel.fa3_utils import page_table_copy
+from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor
+from lightllm.utils.sgl_utils import flash_attn_varlen_func
+
+
+class MlaFa3AttBackend(BaseAttBackend):
+    def __init__(self, model):
+        super().__init__(model=model)
+        self.get_page_table_buffer()  # init
+
+    def get_page_table_buffer(self):
+        """
+        用于减少 decode graph 捕获的时候, 造成显存二次方增长的情况.
+        """
+        model = self.model
+        if not hasattr(self, "_shared_page_table_buffer"):
+            self._shared_page_table_buffer = [
+                torch.empty(model.graph_max_batch_size * model.graph_max_len_in_batch, dtype=torch.int32).to(
+                    get_current_device_id()
+                ),
+                torch.empty(model.graph_max_batch_size * model.graph_max_len_in_batch, dtype=torch.int32).to(
+                    get_current_device_id()
+                ),
+            ]
+        return self._shared_page_table_buffer
+
+    def create_att_prefill_state(self, infer_state) -> "MlaFa3PrefillAttState":
+        return MlaFa3PrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state) -> "MlaFa3DecodeAttState":
+        return MlaFa3DecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class MlaFa3PrefillAttState(BasePrefillAttState):
+    cu_seqlens_q: torch.Tensor = None
+    cu_seqlens_k: torch.Tensor = None
+
+    def init_state(self):
+        self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int()
+        self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int()
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: Tuple[torch.Tensor, torch.Tensor],
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        return self._mla_prefill_att(
+            q=q,
+            k=k,
+            v=v,
+            att_control=att_control,
+            alloc_func=alloc_func,
+        )
+
+    def _mla_prefill_att(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, att_control: AttControl, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        self.backend: MlaFa3AttBackend = self.backend  # for typing
+        k_nope, k_rope = k
+        q_head_num = q.shape[1]
+        k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1)
+
+        assert q.ndim == 3 and k.ndim == 3 and v.ndim == 3
+
+        assert att_control.mla_prefill
+        softmax_scale = att_control.mla_prefill_dict["softmax_scale"]
+
+        o_tensor = flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=self.cu_seqlens_q,
+            cu_seqlens_k=self.cu_seqlens_k,
+            max_seqlen_q=self.infer_state.max_q_seq_len,
+            max_seqlen_k=self.infer_state.max_kv_seq_len,
+            softmax_scale=softmax_scale,
+            causal=True,
+            return_softmax_lse=False,
+        )
+        return o_tensor
+
+
+@dataclasses.dataclass
+class MlaFa3DecodeAttState(BaseDecodeAttState):
+    cu_seqlens_q: torch.Tensor = None
+    cu_seqlens_k: torch.Tensor = None
+    page_table: torch.Tensor = None
+    b_att_seq_len: torch.Tensor = None
+    # 在是否开启mtp 的不同模式下，其设置不同的值，可以加速算子的运行。
+    decode_max_q_seq_len: int = None
+
+    def init_state(self):
+        self.backend: MlaFa3AttBackend = self.backend
+        args_mtp_step = get_env_start_args().mtp_step
+
+        if args_mtp_step > 0:
+            # 修正 mtp 在 fa3 下的输入。
+            mtp_size = args_mtp_step + 1
+            b_q_seq_len = torch.full(
+                (self.infer_state.b_seq_len.shape[0] // mtp_size,),
+                fill_value=mtp_size,
+                dtype=torch.int32,
+                device=self.infer_state.b_seq_len.device,
+            )
+            b_kv_seq_len = self.infer_state.b_seq_len[mtp_size - 1 :: mtp_size]
+            b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len)
+            self.cu_seqlens_q = b1_cu_q_seq_len.int()
+            self.cu_seqlens_k = b1_cu_kv_seq_len.int()
+        else:
+            self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int()
+            self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int()
+
+        att_batch_size = self.infer_state.batch_size // (args_mtp_step + 1)
+        assert self.infer_state.batch_size % (args_mtp_step + 1) == 0
+
+        model = self.backend.model
+        # 可以使用 cuda graph的时候从 buffer中申请
+        if (
+            self.infer_state.batch_size <= model.graph_max_batch_size
+            and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch
+        ):
+            page_buffer = self.backend.get_page_table_buffer()
+            self.page_table = page_buffer[self.infer_state.microbatch_index][
+                : att_batch_size * model.graph_max_len_in_batch
+            ].reshape(att_batch_size, model.graph_max_len_in_batch)
+        else:
+            self.page_table = torch.empty(
+                (att_batch_size, self.infer_state.max_kv_seq_len),
+                dtype=torch.int32,
+                device=self.infer_state.input_ids.device,
+            )
+
+        if args_mtp_step > 0:
+            page_table_copy(
+                page_table=self.page_table[:, : self.infer_state.max_kv_seq_len],
+                req_to_token_indexs=model.req_manager.req_to_token_indexs,
+                b_req_idx=self.infer_state.b_req_idx[args_mtp_step :: (args_mtp_step + 1)],
+            )
+            self.b_att_seq_len = self.infer_state.b_seq_len[args_mtp_step :: (args_mtp_step + 1)].contiguous()
+            self.decode_max_q_seq_len = args_mtp_step + 1
+        else:
+            page_table_copy(
+                page_table=self.page_table[:, : self.infer_state.max_kv_seq_len],
+                req_to_token_indexs=model.req_manager.req_to_token_indexs,
+                b_req_idx=self.infer_state.b_req_idx,
+            )
+            self.b_att_seq_len = self.infer_state.b_seq_len
+            self.decode_max_q_seq_len = 1
+        return
+
+    def copy_for_decode_cuda_graph(self, new_state: "MlaFa3DecodeAttState"):
+        super().copy_for_decode_cuda_graph(new_state)
+
+    def decode_att(
+        self,
+        q: Tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ):
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        assert v is None
+
+        return self._mla_decode_att(
+            q=q,
+            k=k,
+            v=v,
+            att_control=att_control,
+            alloc_func=alloc_func,
+        )
+
+    def _mla_decode_att(
+        self,
+        q: Tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ):
+        q_nope, q_rope = q
+        kv = k
+        qk_rope_head_dim = 64
+        kv_lora_rank = kv.shape[-1] - qk_rope_head_dim
+        k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, 1, 1, qk_rope_head_dim)
+        kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, 1, 1, kv_lora_rank)
+        k_descale, v_descale = None, None
+        assert att_control.mla_decode
+        softmax_scale = att_control.mla_decode_dict["softmax_scale"]
+
+        o_tensor = flash_attn_with_kvcache(
+            q=q_rope,
+            k_cache=k_rope,
+            v_cache=kv_nope,
+            qv=q_nope,
+            page_table=self.page_table,
+            cache_seqlens=self.b_att_seq_len,
+            cu_seqlens_q=self.cu_seqlens_q,
+            cu_seqlens_k_new=self.cu_seqlens_k,
+            max_seqlen_q=self.decode_max_q_seq_len,
+            softmax_scale=softmax_scale,
+            causal=True,
+            window_size=(-1, -1),
+            softcap=0.0,
+            k_descale=k_descale,
+            v_descale=v_descale,
+            return_softmax_lse=False,
+        )
+        return o_tensor
diff --git a/lightllm/common/quantization/triton_quant/__init__.py b/lightllm/common/basemodel/attention/flashinfer/__init__.py
similarity index 100%
rename from lightllm/common/quantization/triton_quant/__init__.py
rename to lightllm/common/basemodel/attention/flashinfer/__init__.py
diff --git a/lightllm/common/basemodel/attention/flashinfer/fp.py b/lightllm/common/basemodel/attention/flashinfer/fp.py
new file mode 100644
index 0000000000..4c6ec0efc6
--- /dev/null
+++ b/lightllm/common/basemodel/attention/flashinfer/fp.py
@@ -0,0 +1,229 @@
+import dataclasses
+import torch
+from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
+from lightllm.utils.dist_utils import get_dp_world_size, get_current_device_id
+from ...triton_kernel.repack_kv_index import repack_kv_index
+
+
+class FlashInferAttBackend(BaseAttBackend):
+    def __init__(self, model):
+        super().__init__(model=model)
+        tp_world_size = get_dp_world_size()
+        self.tp_q_head_num = model.config["num_attention_heads"] // tp_world_size
+        self.tp_kv_head_num = max(model.config["num_key_value_heads"] // tp_world_size, 1)
+        head_dim = model.config["hidden_size"] // model.config["num_attention_heads"]
+        self.head_dim = model.config.get("head_dim", head_dim)
+        self.workspace_buffer = torch.empty(512 * 1024 * 1024, dtype=torch.int8, device=get_current_device_id())
+        self.max_seq_length = model.max_seq_length
+        self.kv_indices_buffer = [
+            torch.empty(
+                model.graph_max_batch_size * self.max_seq_length, dtype=torch.int32, device=get_current_device_id()
+            ),
+            torch.empty(
+                model.graph_max_batch_size * self.max_seq_length, dtype=torch.int32, device=get_current_device_id()
+            ),
+        ]
+        self.q_data_type = model.data_type
+        self.kv_data_type = model.data_type
+
+    def create_att_prefill_state(self, infer_state) -> "FlashInferPrefillAttState":
+        return FlashInferPrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state) -> "FlashInferDecodeAttState":
+        return FlashInferDecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class FlashInferPrefillAttState(BasePrefillAttState):
+    prefill_wrapper: object = None
+
+    def init_state(self):
+        self.backend: FlashInferAttBackend = self.backend
+
+        import flashinfer
+
+        batch_size = self.infer_state.batch_size
+        device = self.infer_state.input_ids.device
+
+        q_starts = self.infer_state.b1_cu_q_seq_len.int()
+        kv_starts = self.infer_state.b1_cu_kv_seq_len.int()
+        kv_last_page_len = torch.full((batch_size,), 1, dtype=torch.int32, device=device)
+        kv_indices = torch.empty(
+            batch_size * self.backend.max_seq_length,
+            dtype=torch.int32,
+            device=device,
+        )
+        repack_kv_index(
+            self.infer_state.req_manager.req_to_token_indexs,
+            self.infer_state.b_req_idx,
+            self.infer_state.b_seq_len,
+            kv_starts[:-1],
+            self.infer_state.max_kv_seq_len,
+            kv_indices,
+        )
+        self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper(
+            self.backend.workspace_buffer,
+            qo_indptr_buf=q_starts,
+            paged_kv_indptr_buf=kv_starts,
+            paged_kv_indices_buf=kv_indices,
+            paged_kv_last_page_len_buf=kv_last_page_len,
+        )
+        self.prefill_wrapper.plan(
+            q_starts,
+            kv_starts,
+            kv_indices,
+            kv_last_page_len,
+            self.backend.tp_q_head_num,
+            self.backend.tp_kv_head_num,
+            self.backend.head_dim,
+            1,
+            causal=True,
+            pos_encoding_mode="NONE",
+            logits_soft_cap=0.0,
+            q_data_type=self.backend.q_data_type,
+            kv_data_type=self.backend.kv_data_type,
+        )
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        return self._nomarl_prefill_att(
+            q=q,
+            k=k,
+            v=v,
+            alloc_func=alloc_func,
+        )
+
+    def _nomarl_prefill_att(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        self.backend: FlashInferAttBackend = self.backend  # for typing
+        o_tensor = alloc_func(q.shape, q.dtype, device="cuda")
+        self.prefill_wrapper.run(
+            q,
+            (k.unsqueeze(1), v.unsqueeze(1)),
+            out=o_tensor,
+        )
+        return o_tensor
+
+
+@dataclasses.dataclass
+class FlashInferDecodeAttState(BaseDecodeAttState):
+    kv_last_page_len_buffer: torch.Tensor = None
+    kv_indices: torch.Tensor = None
+    kv_starts: torch.Tensor = None
+    decode_wrapper: object = None
+
+    def init_state(self):
+        import flashinfer
+
+        self.backend: FlashInferAttBackend = self.backend
+        device = self.infer_state.input_ids.device
+        model = self.backend.model
+        self.kv_last_page_len_buffer = torch.full((self.infer_state.batch_size,), 1, dtype=torch.int32, device=device)
+        if (
+            self.infer_state.batch_size <= model.graph_max_batch_size
+            and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch
+        ):
+            self.kv_indices = self.backend.kv_indices_buffer[self.infer_state.microbatch_index][
+                : self.infer_state.batch_size * self.backend.max_seq_length
+            ]
+        else:
+            self.kv_indices = torch.empty(
+                self.infer_state.batch_size * self.backend.max_seq_length,
+                dtype=torch.int32,
+                device=device,
+            )
+
+        repack_kv_index(
+            self.infer_state.req_manager.req_to_token_indexs,
+            self.infer_state.b_req_idx,
+            self.infer_state.b_seq_len,
+            self.infer_state.b_kv_start_loc,
+            self.infer_state.max_kv_seq_len,
+            self.kv_indices,
+        )
+        self.kv_starts = self.infer_state.b1_cu_kv_seq_len.int()
+        assert self.decode_wrapper is None
+        self.decode_wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(
+            self.backend.workspace_buffer,
+            "NHD",
+            use_cuda_graph=True,
+            use_tensor_cores=True,
+            paged_kv_indptr_buffer=self.kv_starts,
+            paged_kv_indices_buffer=self.kv_indices,
+            paged_kv_last_page_len_buffer=self.kv_last_page_len_buffer,
+        )
+        self.decode_wrapper.plan(
+            self.kv_starts,
+            self.kv_indices,
+            self.kv_last_page_len_buffer,
+            self.backend.tp_q_head_num,
+            self.backend.tp_kv_head_num,
+            self.backend.head_dim,
+            1,
+            q_data_type=self.backend.q_data_type,
+            kv_data_type=self.backend.kv_data_type,
+            non_blocking=True,
+        )
+        return
+
+    def copy_for_decode_cuda_graph(self, new_state: "FlashInferDecodeAttState"):
+        super().copy_for_decode_cuda_graph(new_state)
+        self.decode_wrapper.plan(
+            new_state.kv_starts,
+            new_state.kv_indices,
+            new_state.kv_last_page_len_buffer,
+            new_state.backend.tp_q_head_num,
+            new_state.backend.tp_kv_head_num,
+            new_state.backend.head_dim,
+            1,
+            q_data_type=new_state.backend.q_data_type,
+            kv_data_type=new_state.backend.kv_data_type,
+            non_blocking=True,
+        )
+
+    def decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ):
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        return self._normal_decode_att(
+            q=q,
+            k=k,
+            v=v,
+            alloc_func=alloc_func,
+        )
+
+    def _normal_decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        alloc_func=torch.empty,
+    ):
+        o_tensor = alloc_func(q.shape, q.dtype)
+        self.decode_wrapper.run(
+            q,
+            (k.unsqueeze(1), v.unsqueeze(1)),
+            out=o_tensor,
+        )
+        return o_tensor
diff --git a/lightllm/common/basemodel/attention/flashinfer/fp8.py b/lightllm/common/basemodel/attention/flashinfer/fp8.py
new file mode 100644
index 0000000000..115d6985ac
--- /dev/null
+++ b/lightllm/common/basemodel/attention/flashinfer/fp8.py
@@ -0,0 +1,121 @@
+import dataclasses
+import torch
+from ..base_att import AttControl
+from .fp import FlashInferAttBackend, FlashInferPrefillAttState, FlashInferDecodeAttState
+
+
+class Fp8FlashInferAttBackend(FlashInferAttBackend):
+    def __init__(self, model):
+        super().__init__(model=model)
+        self.kv_data_type = torch.float8_e4m3fn
+
+    def create_att_prefill_state(self, infer_state) -> "Fp8FlashInferPrefillAttState":
+        return Fp8FlashInferPrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state) -> "Fp8FlashInferDecodeAttState":
+        return Fp8FlashInferDecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class Fp8FlashInferPrefillAttState(FlashInferPrefillAttState):
+    offline_scales: torch.Tensor = None
+
+    def init_state(self):
+        super().init_state()
+        self.offline_scales = self.infer_state.mem_manager.scales_list
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        return self._fp8_prefill_att(
+            q=q,
+            k=k,
+            v=v,
+            alloc_func=alloc_func,
+        )
+
+    def _fp8_prefill_att(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        o_tensor = alloc_func(q.shape, q.dtype, device="cuda")
+        k = k.unsqueeze(1).view(torch.float8_e4m3fn)
+        v = v.unsqueeze(1).view(torch.float8_e4m3fn)
+        layer_index = self.backend._find_layer_index(k=k, v=v, att_state=self)
+        offline_scales = self.offline_scales
+        k_descale = offline_scales[layer_index][0] if offline_scales is not None else None
+        v_descale = offline_scales[layer_index][1] if offline_scales is not None else None
+        self.prefill_wrapper.run(
+            q,
+            (k, v),
+            k_scale=k_descale,
+            v_scale=v_descale,
+            out=o_tensor,
+        )
+        return o_tensor
+
+
+@dataclasses.dataclass
+class Fp8FlashInferDecodeAttState(FlashInferDecodeAttState):
+    offline_scales: torch.Tensor = None
+
+    def init_state(self):
+        super().init_state()
+        self.offline_scales = self.infer_state.mem_manager.scales_list
+
+    def copy_for_decode_cuda_graph(self, new_state):
+        return super().copy_for_decode_cuda_graph(new_state)
+
+    def decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ):
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        return self._fp8_decode_att(
+            q=q,
+            k=k,
+            v=v,
+            alloc_func=alloc_func,
+        )
+
+    def _fp8_decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        alloc_func=torch.empty,
+    ):
+        o_tensor = alloc_func(q.shape, q.dtype, device="cuda")
+
+        k = k.unsqueeze(1).view(torch.float8_e4m3fn)
+        v = v.unsqueeze(1).view(torch.float8_e4m3fn)
+        offline_scales = self.offline_scales
+        layer_index = self.backend._find_layer_index(k=k, v=v, att_state=self)
+
+        k_descale = offline_scales[layer_index][0] if offline_scales is not None else None
+        v_descale = offline_scales[layer_index][1] if offline_scales is not None else None
+        self.decode_wrapper.run(
+            q,
+            (k, v),
+            k_scale=k_descale,
+            v_scale=v_descale,
+            out=o_tensor,
+        )
+        return o_tensor
diff --git a/lightllm/common/basemodel/attention/flashinfer/mla.py b/lightllm/common/basemodel/attention/flashinfer/mla.py
new file mode 100644
index 0000000000..6e52203b4f
--- /dev/null
+++ b/lightllm/common/basemodel/attention/flashinfer/mla.py
@@ -0,0 +1,233 @@
+import dataclasses
+import torch
+from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
+from lightllm.utils.dist_utils import get_dp_world_size, get_current_device_id
+from ...triton_kernel.repack_kv_index import repack_kv_index
+from typing import Tuple
+
+
+class MlaFlashInferAttBackend(BaseAttBackend):
+    def __init__(self, model):
+        super().__init__(model=model)
+        num_heads = model.config["num_attention_heads"]
+        self.tp_q_head_num = num_heads // get_dp_world_size()
+        self.qk_nope_head_dim = model.qk_nope_head_dim
+        self.qk_rope_head_dim = model.qk_rope_head_dim
+        self.kv_lora_rank = model.kv_lora_rank
+        self.q_data_type = model.data_type
+        self.kv_data_type = model.data_type
+        self.workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8, device=get_current_device_id())
+        self.max_seq_length = model.max_seq_length
+        self.softmax_scale = (self.qk_nope_head_dim + self.qk_rope_head_dim) ** (-0.5)
+        self.kv_indices_buffer = [
+            torch.empty(
+                model.graph_max_batch_size * self.max_seq_length, dtype=torch.int32, device=get_current_device_id()
+            ),
+            torch.empty(
+                model.graph_max_batch_size * self.max_seq_length, dtype=torch.int32, device=get_current_device_id()
+            ),
+        ]
+
+        from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale
+
+        if model.config["rope_scaling"] is not None:
+            rope_scaling = model.config["rope_scaling"]
+            mscale_all_dim = rope_scaling.get("mscale_all_dim", 0)
+            scaling_factor = rope_scaling["factor"]
+            if mscale_all_dim:
+                mscale = get_deepseek_mscale(scaling_factor, mscale_all_dim)
+                self.softmax_scale = self.softmax_scale * mscale * mscale
+        return
+
+    def create_att_prefill_state(self, infer_state) -> "MlaFlashInferPrefillAttState":
+        return MlaFlashInferPrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state) -> "MlaFlashInferDecodeAttState":
+        return MlaFlashInferDecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class MlaFlashInferPrefillAttState(BasePrefillAttState):
+    prefill_wrapper: object = None
+
+    def init_state(self):
+        self.backend: MlaFlashInferAttBackend = self.backend
+
+        import flashinfer
+
+        q_starts = self.infer_state.b1_cu_q_seq_len.int()
+        kv_starts = self.infer_state.b1_cu_kv_seq_len.int()
+        if self.prefill_wrapper is None:
+            self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper(
+                self.backend.workspace_buffer, "NHD"
+            )
+        self.prefill_wrapper.plan(
+            qo_indptr=q_starts,
+            kv_indptr=kv_starts,
+            num_qo_heads=self.backend.tp_q_head_num,
+            num_kv_heads=self.backend.tp_q_head_num,
+            head_dim_qk=self.backend.qk_nope_head_dim + self.backend.qk_rope_head_dim,
+            head_dim_vo=self.backend.qk_nope_head_dim,
+            q_data_type=self.backend.q_data_type,
+            causal=True,
+            sm_scale=self.backend.softmax_scale,
+        )
+        return
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: Tuple[torch.Tensor, torch.Tensor],
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        return self._mla_prefill_att(
+            q=q,
+            k=k,
+            v=v,
+            alloc_func=alloc_func,
+        )
+
+    def _mla_prefill_att(
+        self, q: torch.Tensor, k: Tuple[torch.Tensor, torch.Tensor], v: torch.Tensor, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        self.backend: MlaFlashInferAttBackend = self.backend  # for typing
+        k_nope, k_rope = k
+        o_tensor = alloc_func((q.shape[0], q.shape[1], k_nope.shape[2]), q.dtype, device="cuda")
+        q_head_num = q.shape[1]
+        k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1)
+        self.prefill_wrapper.run(q, k, v, out=o_tensor)
+        return o_tensor
+
+
+@dataclasses.dataclass
+class MlaFlashInferDecodeAttState(BaseDecodeAttState):
+    kv_indices: torch.Tensor = None
+    kv_starts: torch.Tensor = None
+    decode_wrapper: object = None
+
+    def init_state(self):
+        import flashinfer
+
+        self.backend: MlaFlashInferAttBackend = self.backend
+        model = self.backend.model
+        device = self.infer_state.input_ids.device
+        batch_size = self.infer_state.batch_size
+
+        self.kv_starts = self.infer_state.b1_cu_kv_seq_len
+
+        self.q_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda")
+        if batch_size <= model.graph_max_batch_size and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch:
+            self.kv_indices = self.backend.kv_indices_buffer[self.infer_state.microbatch_index][
+                : batch_size * self.backend.max_seq_length
+            ]
+        else:
+            self.kv_indices = torch.empty(
+                batch_size * self.backend.max_seq_length,
+                dtype=torch.int32,
+                device=device,
+            )
+
+        repack_kv_index(
+            self.infer_state.req_manager.req_to_token_indexs,
+            self.infer_state.b_req_idx,
+            self.infer_state.b_seq_len,
+            self.infer_state.b_kv_start_loc,
+            self.infer_state.max_kv_seq_len,
+            self.kv_indices,
+        )
+        assert self.decode_wrapper is None
+
+        self.decode_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
+            self.backend.workspace_buffer,
+            use_cuda_graph=True,
+            qo_indptr=self.q_indptr,
+            kv_indices=self.kv_indices,
+            kv_indptr=self.kv_starts,
+            kv_len_arr=self.infer_state.b_seq_len,
+        )
+        self.decode_wrapper.plan(
+            self.q_indptr,
+            self.kv_starts,
+            self.kv_indices,
+            self.infer_state.b_seq_len,
+            self.backend.tp_q_head_num,
+            self.backend.kv_lora_rank,
+            self.backend.qk_rope_head_dim,
+            1,
+            False,  # causal
+            self.backend.softmax_scale,
+            self.backend.q_data_type,
+            self.backend.kv_data_type,
+        )
+        return
+
+    def copy_for_decode_cuda_graph(self, new_state: "MlaFlashInferDecodeAttState"):
+        super().copy_for_decode_cuda_graph(new_state)
+        self.decode_wrapper.plan(
+            new_state.q_indptr,
+            new_state.kv_starts,
+            new_state.kv_indices,
+            new_state.infer_state.b_seq_len,
+            new_state.backend.tp_q_head_num,
+            new_state.backend.kv_lora_rank,
+            new_state.backend.qk_rope_head_dim,
+            1,
+            False,  # causal
+            new_state.backend.softmax_scale,
+            new_state.backend.q_data_type,
+            new_state.backend.kv_data_type,
+        )
+
+    def decode_att(
+        self,
+        q: Tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ):
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+
+        assert v is None
+
+        return self._mla_decode_att(
+            q=q,
+            k=k,
+            v=v,
+            att_control=att_control,
+            alloc_func=alloc_func,
+        )
+
+    def _mla_decode_att(
+        self,
+        q: Tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl,
+        alloc_func=torch.empty,
+    ):
+        qk_rope_head_dim = 64
+        q_nope, q_rope = q
+        o_tensor = alloc_func(q_nope.shape, dtype=q_nope.dtype, device=q_nope.device)
+        assert att_control.mla_decode
+
+        self.decode_wrapper.run(
+            q_nope,
+            q_rope,
+            k[:, :, :-qk_rope_head_dim],
+            k[:, :, -qk_rope_head_dim:],
+            out=o_tensor,
+            return_lse=False,
+        )
+        return o_tensor
diff --git a/lightllm/common/quantization/triton_quant/fp8/__init__.py b/lightllm/common/basemodel/attention/triton/__init__.py
similarity index 100%
rename from lightllm/common/quantization/triton_quant/fp8/__init__.py
rename to lightllm/common/basemodel/attention/triton/__init__.py
diff --git a/lightllm/common/basemodel/attention/triton/fp.py b/lightllm/common/basemodel/attention/triton/fp.py
new file mode 100644
index 0000000000..d29f15ec3b
--- /dev/null
+++ b/lightllm/common/basemodel/attention/triton/fp.py
@@ -0,0 +1,275 @@
+import dataclasses
+import torch
+from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
+from typing import Optional
+
+
+class TritonAttBackend(BaseAttBackend):
+    def create_att_prefill_state(self, infer_state) -> "TritonPrefillAttState":
+        return TritonPrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state) -> "TritonDecodeAttState":
+        return TritonDecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class TritonPrefillAttState(BasePrefillAttState):
+    def init_state(self):
+        pass
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert att_control.use_sliding_window is False and att_control.use_att_sink is False
+        if att_control.use_alibi:
+            assert att_control.tp_alibi is not None
+            return self._alibi_prefill_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func)
+        else:
+            return self._nomarl_prefill_att(q=q, k=k, v=v, alloc_func=alloc_func)
+
+    def _alibi_prefill_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl,
+        alloc_func=torch.empty,
+    ):
+        out = alloc_func(q.shape, q.dtype)
+
+        from ...triton_kernel.alibi_att.context_flashattention_nopad import context_attention_fwd
+
+        context_attention_fwd(
+            q,
+            k,
+            v,
+            out,
+            self.infer_state.b_req_idx,
+            att_control.tp_alibi,
+            self.infer_state.b_q_start_loc,
+            self.infer_state.b_seq_len,
+            self.infer_state.b_ready_cache_len,
+            self.infer_state.max_q_seq_len,
+            self.infer_state.req_manager.req_to_token_indexs,
+        )
+        return out
+
+    def _nomarl_prefill_att(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty):
+        from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd
+
+        out = alloc_func(q.shape, q.dtype)
+        context_attention_fwd(
+            q,
+            k,
+            v,
+            out,
+            self.infer_state.b_req_idx,
+            self.infer_state.b_q_start_loc,
+            self.infer_state.b_seq_len,
+            self.infer_state.b_ready_cache_len,
+            self.infer_state.max_q_seq_len,
+            self.infer_state.req_manager.req_to_token_indexs,
+        )
+        return out
+
+
+@dataclasses.dataclass
+class TritonDecodeAttState(BaseDecodeAttState):
+    def init_state(self):
+        pass
+
+    def copy_for_decode_cuda_graph(self, new_state: "TritonDecodeAttState"):
+        super().copy_for_decode_cuda_graph(new_state)
+
+    def decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ):
+        assert att_control.use_sliding_window is False and att_control.use_att_sink is False
+        if att_control.use_alibi:
+            assert att_control.tp_alibi is not None
+            return self._alibi_decode_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func)
+        else:
+            q_head_num = q.shape[1]
+            k_head_num = k.shape[1]
+            if q_head_num == k_head_num:
+                return self._normal_decode_flash_decoding_att(q=q, k=k, v=v, alloc_func=alloc_func)
+            elif q_head_num > k_head_num:
+                return self._normal_decode_gqa_flash_decoding_att(q=q, k=k, v=v, alloc_func=alloc_func)
+            else:
+                raise NotImplementedError("error")
+
+    def _alibi_decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl,
+        alloc_func=torch.empty,
+    ):
+        from ...triton_kernel.alibi_att.token_flashattention_nopad import token_attention_fwd
+
+        out = alloc_func(q.shape, q.dtype)
+        token_attention_fwd(
+            q,
+            k,
+            v,
+            out,
+            att_control.tp_alibi,
+            self.infer_state.req_manager.req_to_token_indexs,
+            self.infer_state.b_req_idx,
+            self.infer_state.b_kv_start_loc,
+            self.infer_state.b_seq_len,
+            self.infer_state.max_kv_seq_len,
+            self.infer_state.total_token_num,
+            alloc_tensor_func=alloc_func,
+        )
+        return out
+
+    def _normal_decode_flash_decoding_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        alloc_func=torch.empty,
+    ):
+        from ...triton_kernel.att.decode_att.mha.flash_decoding.flash_decoding import (
+            token_decode_attention_flash_decoding,
+        )
+
+        out = alloc_func(q.shape, q.dtype)
+
+        token_decode_attention_flash_decoding(
+            q=q,
+            infer_state=self.infer_state,
+            cache_k=k,
+            cache_v=v,
+            out=out,
+            alloc_tensor_func=alloc_func,
+        )
+        return out
+
+    def _normal_decode_gqa_flash_decoding_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        alloc_func=torch.empty,
+    ):
+        from ...triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding import (
+            gqa_token_decode_attention_flash_decoding,
+        )
+
+        out = alloc_func(q.shape, q.dtype)
+
+        gqa_token_decode_attention_flash_decoding(
+            q=q,
+            infer_state=self.infer_state,
+            cache_k=k,
+            cache_v=v,
+            out=out,
+            alloc_tensor_func=alloc_func,
+        )
+
+        return out
+
+    def _normal_decode_gqa_flash_decoding_att_vsm(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        alloc_func=torch.empty,
+    ):
+        # TODO USE , 在特定场景下比 _normal_decode_gqa_flash_decoding_att 省显存
+        from ...triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_vsm import (
+            gqa_token_decode_attention_flash_decoding_vsm,
+        )
+
+        out = alloc_func(q.shape, q.dtype)
+
+        gqa_token_decode_attention_flash_decoding_vsm(
+            q=q,
+            k=k,
+            v=v,
+            infer_state=self.infer_state,
+            out=out,
+            alloc_tensor_func=alloc_func,
+        )
+        return out
+
+    def _normal_decode_gqa_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_weight,
+        alloc_func=torch.empty,
+    ):
+        # TODO USE , 在特定场景下比 _normal_decode_gqa_flash_decoding_att 省显存
+        from ...triton_kernel.att.decode_att.gqa.gqa_decode_flashattention_nopad import gqa_decode_attention_fwd
+
+        out = alloc_func(q.shape, q.dtype)
+
+        gqa_decode_attention_fwd(
+            q=q,
+            k=k,
+            v=v,
+            out=out,
+            req_to_tokens=self.infer_state.req_manager.req_to_token_indexs,
+            b_req_idx=self.infer_state.b_req_idx,
+            b_seq_len=self.infer_state.b_seq_len,
+        )
+        return out
+
+    def _normal_decode_stage3_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        alloc_func=torch.empty,
+    ):
+        total_token_num = self.infer_state.total_token_num
+        batch_size = self.infer_state.batch_size
+        q_head_num = q.shape[1]
+        head_dim = q.shape[2]
+
+        calcu_shape1 = (batch_size, q_head_num, head_dim)
+        att_m_tensor = alloc_func((q_head_num, total_token_num), torch.float32)
+
+        from ...triton_kernel.att.decode_att.mha.stage3_decode_att.token_attention_nopad_att1 import token_att_fwd
+
+        token_att_fwd(
+            q.view(calcu_shape1),
+            k,
+            att_m_tensor,
+            Req_to_tokens=self.infer_state.req_manager.req_to_token_indexs,
+            B_req_idx=self.infer_state.b_req_idx,
+            B_Start_Loc=self.infer_state.b_kv_start_loc,
+            B_Seqlen=self.infer_state.b_seq_len,
+            max_len_in_batch=self.infer_state.max_kv_seq_len,
+        )
+
+        o_tensor = alloc_func(q.shape, q.dtype)
+        from ...triton_kernel.att.decode_att.mha.stage3_decode_att.token_attention_softmax_and_reducev import (
+            token_softmax_reducev_fwd,
+        )
+
+        token_softmax_reducev_fwd(
+            att_m_tensor,
+            v,
+            o_tensor.view(calcu_shape1),
+            req_to_tokens=self.infer_state.req_manager.req_to_token_indexs,
+            b_req_idx=self.infer_state.b_req_idx,
+            b_start_loc=self.infer_state.b_kv_start_loc,
+            b_seq_len=self.infer_state.b_seq_len,
+        )
+        return o_tensor
diff --git a/lightllm/common/basemodel/attention/triton/int4kv.py b/lightllm/common/basemodel/attention/triton/int4kv.py
new file mode 100644
index 0000000000..25199dc470
--- /dev/null
+++ b/lightllm/common/basemodel/attention/triton/int4kv.py
@@ -0,0 +1,170 @@
+import dataclasses
+import torch
+from lightllm.utils.envs_utils import get_env_start_args
+from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
+from typing import Optional, Tuple
+
+
+class Int4kvTritonAttBackend(BaseAttBackend):
+    def __init__(self, model):
+        super().__init__(model)
+        self.quant_group_size: int = get_env_start_args().llm_kv_quant_group_size
+
+    def create_att_prefill_state(self, infer_state) -> "Int4kvTritonPrefillAttState":
+        return Int4kvTritonPrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state) -> "Int4kvTritonDecodeAttState":
+        return Int4kvTritonDecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class Int4kvTritonPrefillAttState(BasePrefillAttState):
+
+    # 用于反量化的时候使用，可以减少反量化占用的显存数量。按需使用。
+    b_kv_start_loc: torch.Tensor = None
+
+    def init_state(self):
+        self.b_kv_start_loc = (
+            torch.cumsum(self.infer_state.b_seq_len, dim=0, dtype=self.infer_state.b_seq_len.dtype)
+            - self.infer_state.b_seq_len
+        )
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: Tuple[torch.Tensor, torch.Tensor],
+        v: Tuple[torch.Tensor, torch.Tensor],
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+
+        self.backend: Int4kvTritonAttBackend = self.backend  # for typing
+
+        k, k_scale = k
+        v, v_scale = v
+        o = self._groupsize_quant_prefill_att(
+            q=q,
+            k=k,
+            k_scale=k_scale,
+            v=v,
+            v_scale=v_scale,
+            alloc_func=alloc_func,
+        )
+        return o
+
+    def _groupsize_quant_prefill_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        k_scale: torch.Tensor,
+        v: torch.Tensor,
+        v_scale: torch.Tensor,
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        # o_tensor = alloc_func(q.shape, q.dtype, device=q.device)
+        # batch_size = self.infer_state.b_seq_len.shape[0]
+
+        assert k.untyped_storage().data_ptr() == v.untyped_storage().data_ptr()
+        assert k_scale.untyped_storage().data_ptr() == v_scale.untyped_storage().data_ptr()
+
+        total_token_num = self.infer_state.total_token_num
+        head_dim = k.shape[2] * 2  # 2个4bit存储为一个int8, 所以维度需要翻倍，才是解量化后的精度
+        k_dequant = alloc_func((total_token_num, k.shape[1], head_dim), dtype=q.dtype, device=q.device)
+        v_dequant = alloc_func((total_token_num, v.shape[1], head_dim), dtype=q.dtype, device=q.device)
+        o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device)
+
+        max_kv_seq_len = self.infer_state.max_kv_seq_len
+
+        from ...triton_kernel.kv_copy.ppl_int4kv_copy_kv import dequantize_int4kv
+
+        dequantize_int4kv(
+            k=k,
+            k_scale=k_scale,
+            v=v,
+            v_scale=v_scale,
+            req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs,
+            b_seq_len=self.infer_state.b_seq_len,
+            b_req_idx=self.infer_state.b_req_idx,
+            b_kv_start_loc=self.b_kv_start_loc,
+            k_out=k_dequant,
+            v_out=v_dequant,
+            max_len_in_batch=max_kv_seq_len,
+            quant_group_size=self.backend.quant_group_size,
+        )
+
+        from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd_contiguous_kv
+
+        context_attention_fwd_contiguous_kv(
+            q=q,
+            k=k_dequant,
+            v=v_dequant,
+            o=o_tensor,
+            b_start_loc=self.infer_state.b_q_start_loc,
+            b_kv_start_loc=self.b_kv_start_loc,
+            b_seq_len=self.infer_state.b_seq_len,
+            max_q_input_len=self.infer_state.max_q_seq_len,
+            b_prompt_cache_len=self.infer_state.b_ready_cache_len,
+        )
+        return o_tensor
+
+
+@dataclasses.dataclass
+class Int4kvTritonDecodeAttState(BaseDecodeAttState):
+    def init_state(self):
+        pass
+
+    def copy_for_decode_cuda_graph(self, new_state: "Int4kvTritonDecodeAttState"):
+        super().copy_for_decode_cuda_graph(new_state)
+
+    def decode_att(
+        self,
+        q: torch.Tensor,
+        k: Tuple[torch.Tensor, torch.Tensor],
+        v: Tuple[torch.Tensor, torch.Tensor],
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ):
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        k, k_scale = k
+        v, v_scale = v
+
+        return self.ppl_int4kv_decode_att(
+            q=q,
+            k=k,
+            k_scale=k_scale,
+            v=v,
+            v_scale=v_scale,
+            alloc_func=alloc_func,
+        )
+
+    def ppl_int4kv_decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        k_scale: torch.Tensor,
+        v: torch.Tensor,
+        v_scale: torch.Tensor,
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        from ...triton_kernel.att.decode_att.int4kv.ppl_int4kv_flash_decoding import (
+            token_decode_attention_flash_decoding,
+        )
+
+        return token_decode_attention_flash_decoding(
+            q=q,
+            infer_state=self.infer_state,
+            cache_k=k,
+            cache_k_scale=k_scale,
+            cache_v=v,
+            cache_v_scale=v_scale,
+            alloc_tensor_func=alloc_func,
+        )
diff --git a/lightllm/common/basemodel/attention/triton/int8kv.py b/lightllm/common/basemodel/attention/triton/int8kv.py
new file mode 100644
index 0000000000..975d7b629c
--- /dev/null
+++ b/lightllm/common/basemodel/attention/triton/int8kv.py
@@ -0,0 +1,196 @@
+import dataclasses
+import torch
+from lightllm.utils.envs_utils import get_env_start_args
+from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
+from typing import Optional, Tuple
+from lightllm.utils.envs_utils import enable_diverse_mode_gqa_decode_fast_kernel
+
+
+class Int8kvTritonAttBackend(BaseAttBackend):
+    def __init__(self, model):
+        super().__init__(model)
+        self.quant_group_size: int = get_env_start_args().llm_kv_quant_group_size
+
+    def create_att_prefill_state(self, infer_state) -> "Int8kvTritonPrefillAttState":
+        return Int8kvTritonPrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state) -> "Int8kvTritonDecodeAttState":
+        return Int8kvTritonDecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class Int8kvTritonPrefillAttState(BasePrefillAttState):
+
+    # 用于反量化的时候使用，可以减少反量化占用的显存数量。按需使用。
+    b_kv_start_loc: torch.Tensor = None
+
+    def init_state(self):
+        self.b_kv_start_loc = (
+            torch.cumsum(self.infer_state.b_seq_len, dim=0, dtype=self.infer_state.b_seq_len.dtype)
+            - self.infer_state.b_seq_len
+        )
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: Tuple[torch.Tensor, torch.Tensor],
+        v: Tuple[torch.Tensor, torch.Tensor],
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+
+        self.backend: Int8kvTritonAttBackend = self.backend  # for typing
+
+        k, k_scale = k
+        v, v_scale = v
+        o = self._groupsize_quant_prefill_att(
+            q=q,
+            k=k,
+            k_scale=k_scale,
+            v=v,
+            v_scale=v_scale,
+            alloc_func=alloc_func,
+        )
+        return o
+
+    def _groupsize_quant_prefill_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        k_scale: torch.Tensor,
+        v: torch.Tensor,
+        v_scale: torch.Tensor,
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        # o_tensor = alloc_func(q.shape, q.dtype, device=q.device)
+        # batch_size = self.infer_state.b_seq_len.shape[0]
+
+        assert k.untyped_storage().data_ptr() == v.untyped_storage().data_ptr()
+        assert k_scale.untyped_storage().data_ptr() == v_scale.untyped_storage().data_ptr()
+
+        total_token_num = self.infer_state.total_token_num
+        k_dequant = alloc_func((total_token_num, k.shape[1], k.shape[2]), dtype=q.dtype, device=q.device)
+        v_dequant = alloc_func((total_token_num, v.shape[1], v.shape[2]), dtype=q.dtype, device=q.device)
+        o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device)
+
+        max_kv_seq_len = self.infer_state.max_kv_seq_len
+
+        from ...triton_kernel.kv_copy.ppl_int8kv_copy_kv import dequantize_int8kv
+
+        dequantize_int8kv(
+            k=k,
+            k_scale=k_scale,
+            v=v,
+            v_scale=v_scale,
+            req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs,
+            b_seq_len=self.infer_state.b_seq_len,
+            b_req_idx=self.infer_state.b_req_idx,
+            b_kv_start_loc=self.b_kv_start_loc,
+            k_out=k_dequant,
+            v_out=v_dequant,
+            max_len_in_batch=max_kv_seq_len,
+            quant_group_size=self.backend.quant_group_size,
+        )
+
+        from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd_contiguous_kv
+
+        context_attention_fwd_contiguous_kv(
+            q=q,
+            k=k_dequant,
+            v=v_dequant,
+            o=o_tensor,
+            b_start_loc=self.infer_state.b_q_start_loc,
+            b_kv_start_loc=self.b_kv_start_loc,
+            b_seq_len=self.infer_state.b_seq_len,
+            max_q_input_len=self.infer_state.max_q_seq_len,
+            b_prompt_cache_len=self.infer_state.b_ready_cache_len,
+        )
+        return o_tensor
+
+
+@dataclasses.dataclass
+class Int8kvTritonDecodeAttState(BaseDecodeAttState):
+    def init_state(self):
+        pass
+
+    def copy_for_decode_cuda_graph(self, new_state: "Int8kvTritonDecodeAttState"):
+        super().copy_for_decode_cuda_graph(new_state)
+
+    def decode_att(
+        self,
+        q: torch.Tensor,
+        k: Tuple[torch.Tensor, torch.Tensor],
+        v: Tuple[torch.Tensor, torch.Tensor],
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ):
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        k, k_scale = k
+        v, v_scale = v
+        if enable_diverse_mode_gqa_decode_fast_kernel():
+            return self.diverse_decode_att(q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func)
+        else:
+            return self.ppl_mha_int8kv_decode_att(
+                q=q,
+                k=k,
+                k_scale=k_scale,
+                v=v,
+                v_scale=v_scale,
+                alloc_func=alloc_func,
+            )
+
+    def diverse_decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        k_scale: torch.Tensor,
+        v: torch.Tensor,
+        v_scale: torch.Tensor,
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+
+        from ...triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse import (
+            token_decode_attention_flash_decoding,
+        )
+
+        return token_decode_attention_flash_decoding(
+            q=q,
+            infer_state=self.infer_state,
+            cache_k=k,
+            cache_k_scale=k_scale,
+            cache_v=v,
+            cache_v_scale=v_scale,
+            alloc_tensor_func=alloc_func,
+        )
+
+    def ppl_mha_int8kv_decode_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        k_scale: torch.Tensor,
+        v: torch.Tensor,
+        v_scale: torch.Tensor,
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        from ...triton_kernel.att.decode_att.int8kv.ppl_int8kv_flash_decoding import (
+            token_decode_attention_flash_decoding,
+        )
+
+        return token_decode_attention_flash_decoding(
+            q=q,
+            infer_state=self.infer_state,
+            cache_k=k,
+            cache_k_scale=k_scale,
+            cache_v=v,
+            cache_v_scale=v_scale,
+            alloc_tensor_func=alloc_func,
+        )
diff --git a/lightllm/common/basemodel/attention/triton/mla.py b/lightllm/common/basemodel/attention/triton/mla.py
new file mode 100644
index 0000000000..8288193ad7
--- /dev/null
+++ b/lightllm/common/basemodel/attention/triton/mla.py
@@ -0,0 +1,125 @@
+import dataclasses
+import torch
+from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
+from typing import Tuple
+
+
+class MlaTritonAttBackend(BaseAttBackend):
+    def create_att_prefill_state(self, infer_state) -> "MlaTritonPrefillAttState":
+        return MlaTritonPrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state) -> "MlaTritonDecodeAttState":
+        return MlaTritonDecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class MlaTritonPrefillAttState(BasePrefillAttState):
+    def init_state(self):
+        pass
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: Tuple[torch.Tensor, torch.Tensor],
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert (
+            att_control.use_alibi is False
+            and att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+        )
+        return self._mla_prefill_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func)
+
+    def _mla_prefill_att(
+        self,
+        q: torch.Tensor,
+        k: Tuple[torch.Tensor, torch.Tensor],
+        v: torch.Tensor,
+        att_control: AttControl,
+        alloc_func=torch.empty,
+    ):
+        from ...triton_kernel.mla_att.prefill_att import context_attention_fwd_with_v
+
+        qk_rope_head_dim = 64
+        q_nope, q_rope = q[:, :, :-qk_rope_head_dim], q[:, :, -qk_rope_head_dim:]
+        o_tensor = alloc_func(q_nope.shape, dtype=q_nope.dtype, device=q.device)
+        k_nope, k_rope = k
+        assert att_control.mla_prefill
+        softmax_scale = att_control.mla_prefill_dict["softmax_scale"]
+        context_attention_fwd_with_v(
+            q_nope,
+            q_rope,
+            k_nope,
+            k_rope,
+            v,
+            o_tensor,
+            self.infer_state.b_q_start_loc,
+            self.infer_state.b1_cu_kv_seq_len,
+            self.infer_state.b_seq_len,
+            self.infer_state.b_ready_cache_len,
+            self.infer_state.max_q_seq_len,
+            softmax_scale,
+        )
+        return o_tensor
+
+
+@dataclasses.dataclass
+class MlaTritonDecodeAttState(BaseDecodeAttState):
+    def init_state(self):
+        pass
+
+    def copy_for_decode_cuda_graph(self, new_state: "MlaTritonDecodeAttState"):
+        super().copy_for_decode_cuda_graph(new_state)
+
+    def decode_att(
+        self,
+        q: Tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ):
+        assert (
+            att_control.use_sliding_window is False
+            and att_control.use_att_sink is False
+            and att_control.use_alibi is False
+        )
+        assert v is None
+
+        return self._mla_decode_att(
+            q=q,
+            k=k,
+            v=v,
+            att_control=att_control,
+            alloc_func=alloc_func,
+        )
+
+    def _mla_decode_att(
+        self,
+        q: Tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl,
+        alloc_func=torch.empty,
+    ):
+        assert att_control.mla_decode
+        softmax_scale = att_control.mla_decode_dict["softmax_scale"]
+
+        from ...triton_kernel.mla_att.decode_att import gqa_token_decode_attention_flash_decoding
+
+        qk_rope_head_dim = 64
+        q_nope, q_rope = q
+        kv = k
+
+        out = gqa_token_decode_attention_flash_decoding(
+            q_nope=q_nope,
+            q_rope=q_rope,
+            kv_nope=kv[:, :, :-qk_rope_head_dim],
+            kv_rope=kv[:, :, -qk_rope_head_dim:],
+            infer_state=self.infer_state,
+            softmax_scale=softmax_scale,
+            alloc_tensor_func=alloc_func,
+        )
+        return out
diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py
index 011f998fc0..e6405e4d74 100755
--- a/lightllm/common/basemodel/basemodel.py
+++ b/lightllm/common/basemodel/basemodel.py
@@ -32,6 +32,8 @@
 from lightllm.utils.envs_utils import set_model_init_status, enable_diverse_mode_gqa_decode_fast_kernel
 from lightllm.common.triton_utils.autotuner import Autotuner
 from lightllm.utils.infer_utils import post_empty_cache
+from .attention import get_prefill_att_backend_class, get_decode_att_backend_class
+from .attention import BaseAttBackend
 
 logger = init_logger(__name__)
 
@@ -58,7 +60,6 @@ def __init__(self, kvargs):
         self.max_total_token_num = kvargs["max_total_token_num"]
         self.batch_max_tokens = kvargs.get("batch_max_tokens", None)
         self.load_way = kvargs.get("load_way", "HF")
-        self.mode = kvargs.get("mode", [])
         self.weight_dict = kvargs.get("weight_dict", None)
         self.finetune_config = kvargs.get("finetune_config", None)
         self.max_req_num = kvargs.get("max_req_num", 1000)
@@ -102,23 +103,27 @@ def __init__(self, kvargs):
         self._verify_params()
         self._init_quant()
 
-        # 更连续的显存分配可以有更好的性能
-        if self.max_total_token_num is None:
-            self._init_weights()
-            self._init_mem_manager()
-        else:
-            self._init_mem_manager()
-            self._init_weights()
-
+        self._init_weights()
+        self._init_mem_manager()
         self._init_kv_move_buffer()
         self._check_mem_size()
         self._init_req_manager()
         self._init_infer_layer()
         self._init_some_value()
         self._init_custom()
-        self._init_inferstate_cls()
+        self._load_hf_weights()
         # wait必须在init cudagraph 之前，避免错误捕获
         self._wait_other_modules_ready()
+
+        self._init_att_backend()
+        self._init_att_backend1()
+
+        logger.info(f"use prefill att backend: {self.prefill_att_backend.__class__.__name__}")
+        logger.info(f"use decode att backend: {self.decode_att_backend.__class__.__name__}")
+        if self.prefill_att_backend1 is not None:
+            logger.info(f"use prefill att backend1: {self.prefill_att_backend1.__class__.__name__}")
+            logger.info(f"use decode att backend1: {self.decode_att_backend1.__class__.__name__}")
+
         self._autotune_warmup()
         self._init_padded_req()
         self._init_cudagraph()
@@ -144,9 +149,6 @@ def _init_config(self):
             self.config["vocab_size"] = self.finetune_config.vocab_size
         return
 
-    def _init_inferstate_cls(self):
-        pass
-
     @final
     def _verify_must(self):
         assert self.config["num_attention_heads"] % self.tp_world_size_ == 0
@@ -162,19 +164,19 @@ def _init_quant(self):
         logger.info(f"Initial quantization. " f"The default quantization method is {self.quant_cfg.quant_type}")
 
     def _init_weights(self, start_layer_index=0):
-        self.pre_post_weight = self.pre_and_post_weight_class(
-            self.data_type, network_config=self.config, mode=self.mode
-        )
+        self.pre_post_weight = self.pre_and_post_weight_class(self.data_type, network_config=self.config)
         self.trans_layers_weight = [
             self.transformer_weight_class(
                 i,
                 self.data_type,
                 network_config=self.config,
-                mode=self.mode,
                 quant_cfg=self.quant_cfg,
             )
             for i in range(start_layer_index, start_layer_index + self.config["n_layer"])
         ]
+        return
+
+    def _load_hf_weights(self):
         load_hf_weights(
             self.data_type,
             weight_dir=self.weight_dir_,
@@ -220,10 +222,10 @@ def _init_req_manager(self):
         return
 
     def _init_infer_layer(self, start_layer_index=0):
-        self.pre_infer = self.pre_layer_infer_class(network_config=self.config, mode=self.mode)
-        self.post_infer = self.post_layer_infer_class(network_config=self.config, mode=self.mode)
+        self.pre_infer = self.pre_layer_infer_class(network_config=self.config)
+        self.post_infer = self.post_layer_infer_class(network_config=self.config)
         self.layers_infer = [
-            self.transformer_layer_infer_class(i, network_config=self.config, mode=self.mode)
+            self.transformer_layer_infer_class(i, network_config=self.config)
             for i in range(start_layer_index, start_layer_index + self.config["n_layer"])
         ]
         return
@@ -238,6 +240,17 @@ def _init_some_value(self):
         self.vocab_size = self.config["vocab_size"]
         return
 
+    def _init_att_backend(self):
+        self.prefill_att_backend: BaseAttBackend = get_prefill_att_backend_class(index=0)(model=self)
+        self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0)(model=self)
+        return
+
+    def _init_att_backend1(self):
+        # self.prefill_att_backend1 是给后续有模型支持不同层用不同的att模块时，保留的扩展。
+        self.prefill_att_backend1: BaseAttBackend = None
+        self.decode_att_backend1: BaseAttBackend = None
+        return
+
     def _init_cudagraph(self):
         self.graph = (
             None if self.disable_cudagraph else CudaGraph(self.graph_max_batch_size, self.graph_max_len_in_batch)
@@ -281,7 +294,6 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0)
         infer_state.return_all_prompt_logics = self.return_all_prompt_logics
         infer_state.batch_size = model_input.batch_size
         infer_state.total_token_num = model_input.total_token_num
-        infer_state.max_len_in_batch = model_input.max_len_in_batch
         infer_state.max_q_seq_len = model_input.max_q_seq_len
         infer_state.max_kv_seq_len = model_input.max_kv_seq_len
         infer_state.max_cache_len = model_input.max_cache_len
@@ -311,6 +323,19 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0)
         # 特殊模型，特殊模式的特定变量初始化操作。
         infer_state.mtp_draft_input_hiddens = model_input.mtp_draft_input_hiddens
 
+        if infer_state.is_prefill:
+            infer_state.prefill_att_state = self.prefill_att_backend.create_att_prefill_state(infer_state=infer_state)
+            if self.prefill_att_backend1 is not None:
+                infer_state.prefill_att_state1 = self.prefill_att_backend1.create_att_prefill_state(
+                    infer_state=infer_state
+                )
+        else:
+            infer_state.decode_att_state = self.decode_att_backend.create_att_decode_state(infer_state=infer_state)
+            if self.decode_att_backend1 is not None:
+                infer_state.decode_att_state1 = self.decode_att_backend1.create_att_decode_state(
+                    infer_state=infer_state
+                )
+
         return infer_state
 
     def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_size: int):
@@ -323,6 +348,7 @@ def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_s
         new_model_input = copy.copy(model_input)
         new_model_input.batch_size = new_batch_size
         new_model_input.total_token_num += padded_batch_size * 2
+        new_model_input.max_kv_seq_len = max(2, model_input.max_kv_seq_len)
         new_model_input.input_ids = F.pad(new_model_input.input_ids, (0, padded_batch_size), mode="constant", value=1)
         new_model_input.b_req_idx = F.pad(
             new_model_input.b_req_idx, (0, padded_batch_size), mode="constant", value=self.req_manager.HOLD_REQUEST_ID
@@ -366,7 +392,6 @@ def _create_padded_prefill_model_input(self, model_input: ModelInput, new_handle
         new_model_input = copy.copy(model_input)
         new_model_input.batch_size = model_input.batch_size + 1
         new_model_input.total_token_num += padded_token_num
-        new_model_input.max_len_in_batch = max(padded_token_num, model_input.max_len_in_batch)
         new_model_input.max_q_seq_len = max(padded_token_num, model_input.max_q_seq_len)
         new_model_input.max_kv_seq_len = max(padded_token_num, model_input.max_kv_seq_len)
         new_model_input.max_cache_len = max(0, model_input.max_cache_len)
@@ -464,6 +489,7 @@ def _prefill(
         prefill_mem_indexes_ready_event.record()
 
         infer_state.init_some_extra_state(self)
+        infer_state.init_att_state()
         model_output = self._context_forward(infer_state)
         if is_padded_model_input:
             model_output = self._create_unpad_prefill_model_output(
@@ -484,7 +510,7 @@ def _decode(
                 model_input.b_mtp_index,
             )
 
-        if self.graph is not None and self.graph.can_run(model_input.batch_size, model_input.max_len_in_batch):
+        if self.graph is not None and self.graph.can_run(model_input.batch_size, model_input.max_kv_seq_len):
             find_graph_batch_size = self.graph.find_closest_graph_batch_size(model_input.batch_size)
             padded_model_input = self._create_padded_decode_model_input(model_input, find_graph_batch_size)
             infer_state = self._create_inferstate(padded_model_input)
@@ -495,6 +521,7 @@ def _decode(
                 infer_state.mem_index,
             )
             infer_state.init_some_extra_state(self)
+            infer_state.init_att_state()
 
             if self.graph.need_capture(find_graph_batch_size):
                 infer_state.is_cuda_graph = True
@@ -514,6 +541,7 @@ def _decode(
                 infer_state.mem_index,
             )
             infer_state.init_some_extra_state(self)
+            infer_state.init_att_state()
             model_output = self._token_forward(infer_state)
 
         return model_output
@@ -609,8 +637,6 @@ def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: Mod
 
         assert model_input0.mem_indexes.is_cuda
         assert model_input1.mem_indexes.is_cuda
-        input_ids0, input_ids1 = model_input0.input_ids, model_input1.input_ids
-
         infer_state0 = self._create_inferstate(model_input0, 0)
         init_req_to_token_indexes(
             req_to_token_indexs=self.req_manager.req_to_token_indexs,
@@ -622,6 +648,7 @@ def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: Mod
             max_q_seq_len=infer_state0.max_q_seq_len,
         )
         infer_state0.init_some_extra_state(self)
+        infer_state0.init_att_state()
 
         infer_state1 = self._create_inferstate(model_input1, 1)
         init_req_to_token_indexes(
@@ -634,13 +661,12 @@ def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: Mod
             max_q_seq_len=infer_state1.max_q_seq_len,
         )
         infer_state1.init_some_extra_state(self)
+        infer_state1.init_att_state()
 
         prefill_mem_indexes_ready_event = torch.cuda.Event()
         prefill_mem_indexes_ready_event.record()
 
-        model_output0, model_output1 = self._overlap_tpsp_context_forward(
-            input_ids0, infer_state0, input_ids1=input_ids1, infer_state1=infer_state1
-        )
+        model_output0, model_output1 = self._overlap_tpsp_context_forward(infer_state0, infer_state1=infer_state1)
 
         # 在开启使用deepep的时候，需要调用clear_deepep_buffer做资源清理，没有启用的时候
         # 该调用没有实际意义
@@ -672,7 +698,7 @@ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: Mode
         assert model_input1.mem_indexes.is_cuda
 
         origin_batch_size = model_input0.batch_size
-        max_len_in_batch = max(model_input0.max_len_in_batch, model_input1.max_len_in_batch)
+        max_len_in_batch = max(model_input0.max_kv_seq_len, model_input1.max_kv_seq_len)
 
         if self.graph is not None and self.graph.can_run(origin_batch_size, max_len_in_batch):
             find_graph_batch_size = self.graph.find_closest_graph_batch_size(origin_batch_size)
@@ -688,6 +714,8 @@ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: Mode
                 infer_state0.mem_index,
             )
             infer_state0.init_some_extra_state(self)
+            infer_state0.init_att_state()
+
             infer_state1 = self._create_inferstate(padded_model_input1, 1)
             copy_kv_index_to_req(
                 self.req_manager.req_to_token_indexs,
@@ -696,6 +724,7 @@ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: Mode
                 infer_state1.mem_index,
             )
             infer_state1.init_some_extra_state(self)
+            infer_state1.init_att_state()
 
             if self.graph.need_capture(find_graph_batch_size):
                 infer_state0.is_cuda_graph = True
@@ -724,6 +753,8 @@ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: Mode
                 infer_state0.mem_index,
             )
             infer_state0.init_some_extra_state(self)
+            infer_state0.init_att_state()
+
             infer_state1 = self._create_inferstate(model_input1, 1)
             copy_kv_index_to_req(
                 self.req_manager.req_to_token_indexs,
@@ -732,6 +763,7 @@ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: Mode
                 infer_state1.mem_index,
             )
             infer_state1.init_some_extra_state(self)
+            infer_state1.init_att_state()
 
             model_output0, model_output1 = self._overlap_tpsp_token_forward(infer_state0, infer_state1=infer_state1)
         return model_output0, model_output1
@@ -818,7 +850,6 @@ def _check_max_len_infer(self):
             model_input = ModelInput(
                 batch_size=1,
                 total_token_num=total_token_num,
-                max_len_in_batch=self.batch_max_tokens,
                 max_q_seq_len=self.batch_max_tokens,
                 max_kv_seq_len=self.batch_max_tokens,
                 max_cache_len=0,
@@ -895,7 +926,6 @@ def _autotune_warmup(self):
                 model_input = ModelInput(
                     batch_size=1,
                     total_token_num=total_token_num,
-                    max_len_in_batch=input_len,
                     max_q_seq_len=input_len,
                     max_kv_seq_len=input_len,
                     max_cache_len=0,
@@ -958,7 +988,6 @@ def _init_padded_req(self):
         model_input = ModelInput(
             batch_size=batch_size,
             total_token_num=total_token_num,
-            max_len_in_batch=prefill_input_len,
             max_q_seq_len=prefill_input_len,
             max_kv_seq_len=prefill_input_len,
             max_cache_len=0,
diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py
index 138f084270..758c0b5194 100644
--- a/lightllm/common/basemodel/batch_objs.py
+++ b/lightllm/common/basemodel/batch_objs.py
@@ -11,10 +11,7 @@ class ModelInput:
     # 通用变量
     batch_size: int
     total_token_num: int
-    max_len_in_batch: int
-    # 在 decode 阶段， 常规模式下， max_q_seq_len 必定是 1，
-    # 在 mtp 模式下，max_q_seq_len 统计的是一个请求考虑了 mtp 步数的
-    # 最大长度，实际值是 max([(1 + req.mtp_step) for req in reqs])
+    # 在 decode 阶段， max_q_seq_len 必定是 1，
     max_q_seq_len: int
     max_kv_seq_len: int
     max_cache_len: int = None
diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py
index 9eeab7270c..dd29c9a833 100644
--- a/lightllm/common/basemodel/cuda_graph.py
+++ b/lightllm/common/basemodel/cuda_graph.py
@@ -67,7 +67,7 @@ def _capture_decode(self, decode_func, infer_state: InferStateInfo):
         graph_obj = torch.cuda.CUDAGraph()
         input_ids = infer_state.input_ids
         batch_size = input_ids.shape[0]
-        infer_state.max_len_in_batch = self.graph_max_len_in_batch
+        infer_state.max_kv_seq_len = self.graph_max_len_in_batch
         infer_state.total_token_num = self.graph_max_len_in_batch * batch_size
         # warmup
         # 因为有些推理过程的代码，会通过判断infer_state中是否存在某些属性来在一层上
@@ -77,10 +77,16 @@ def _capture_decode(self, decode_func, infer_state: InferStateInfo):
         # 浅拷贝，不然后续传入到cuda graph捕获过程中后，infer_state因为提前拥有了这些属性，
         # 导致不会重新初始化，这样捕获过程中会不能捕获这些临时添加到 infer_state 管理对象
         # 中的 tensor。
+
         for _ in range(1):
+            # 记录原始存在的变量
+            pure_para_set = set(vars(infer_state).keys())
             torch.cuda.synchronize()
             decode_func(copy.copy(infer_state))
             torch.cuda.synchronize()
+            for param_name in set(vars(infer_state).keys()):
+                if param_name not in pure_para_set:
+                    delattr(infer_state, param_name)
 
         with lightllm_capture_graph(dist_group):
             with torch.cuda.graph(graph_obj, pool=self.mempool):
@@ -100,15 +106,25 @@ def _capture_decode_overlap(
         graph_obj = torch.cuda.CUDAGraph()
         input_ids = infer_state.input_ids
         batch_size = input_ids.shape[0]
-        infer_state.max_len_in_batch = self.graph_max_len_in_batch
+        infer_state.max_kv_seq_len = self.graph_max_len_in_batch
         infer_state.total_token_num = self.graph_max_len_in_batch * batch_size
-        infer_state1.max_len_in_batch = self.graph_max_len_in_batch
+        infer_state1.max_kv_seq_len = self.graph_max_len_in_batch
         infer_state1.total_token_num = self.graph_max_len_in_batch * batch_size
         # warmup
         for _ in range(1):
+            # 记录原始存在的变量
+            pure_para_set = set(vars(infer_state).keys())
+            pure_para_set1 = set(vars(infer_state1).keys())
             torch.cuda.synchronize()
             decode_func(copy.copy(infer_state), copy.copy(infer_state1))
             torch.cuda.synchronize()
+            for para_name in set(vars(infer_state).keys()):
+                if para_name not in pure_para_set:
+                    delattr(infer_state, para_name)
+            for para_name in set(vars(infer_state1).keys()):
+                if para_name not in pure_para_set1:
+                    delattr(infer_state1, para_name)
+
         with lightllm_capture_graph(dist_group1):
             with lightllm_capture_graph(dist_group):
                 with torch.cuda.graph(graph_obj, pool=self.mempool):
@@ -196,8 +212,7 @@ def warmup(self, model):
             model_input = ModelInput(
                 batch_size=batch_size,
                 total_token_num=total_token_num,
-                max_len_in_batch=max_len_in_batch,
-                max_q_seq_len=self.mtp_step + 1,
+                max_q_seq_len=1,
                 max_kv_seq_len=max_len_in_batch,
                 input_ids=input_ids,
                 mem_indexes=mem_indexes,
@@ -256,8 +271,7 @@ def warmup_overlap(self, model):
                     is_prefill=False,
                     batch_size=batch_size,
                     total_token_num=total_token_num,
-                    max_len_in_batch=max_len_in_batch,
-                    max_q_seq_len=self.mtp_step + 1,
+                    max_q_seq_len=1,
                     max_kv_seq_len=max_len_in_batch,
                     input_ids=input_ids,
                     b_mtp_index=b_mtp_index,
diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py
index 8e7174bb39..75856b1086 100755
--- a/lightllm/common/basemodel/infer_struct.py
+++ b/lightllm/common/basemodel/infer_struct.py
@@ -11,6 +11,7 @@
 from .batch_objs import ModelInput
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.dist_utils import get_global_dp_rank
+from .attention import BasePrefillAttState, BaseDecodeAttState
 
 
 class InferStateInfo:
@@ -19,21 +20,24 @@ class InferStateInfo:
     """
 
     def __init__(self):
+        # prefill 和 decode 使用的 att 状态对象
+        self.prefill_att_state: BasePrefillAttState = None
+        self.decode_att_state: BaseDecodeAttState = None
+
+        # 保留的扩展, 支持线性att与标准att混合使用时使用
+        self.prefill_att_state1: BasePrefillAttState = None
+        self.decode_att_state1: BaseDecodeAttState = None
+
         self.input_ids: torch.Tensor = None
         self.batch_size: int = None
         self.total_token_num: int = None
         self.b_req_idx: torch.Tensor = None
-        self.b_start_loc: torch.Tensor = None
         self.b_ready_cache_len: torch.Tensor = None  # only for prefill prompt cache used.
 
         self.b_shared_seq_len: torch.Tensor = None  # only for diverse mode used in decode phase.
         self.b_mark_shared_group: torch.Tensor = None  # only for diverse mode used in decode phase.
 
         self.b_seq_len: torch.Tensor = None
-        # max_len_in_batch prefill 和 decode 阶段含义不同
-        # prefill 阶段指每个req 输入token的长度（不包括已经cache的部分）最大值
-        # decode 阶段指的是每个req的总长 最大值
-        self.max_len_in_batch: int = None
         # max_cache_len 用于 prefill 阶段标识请求中最大 cache的kv 的长度
         self.max_cache_len: int = None
         # prefix_total_token_num 用于 prefill 阶段标识当前请求中所有已经ready的kv的长度
@@ -68,6 +72,11 @@ def __init__(self):
         self.max_q_seq_len: int = None
         self.max_kv_seq_len: int = None
 
+        # prefill 用
+        self.b_q_start_loc: torch.Tensor = None
+        # decode 用
+        self.b_kv_start_loc: torch.Tensor = None
+
         # 一些特殊模型，特殊模式使用的输入变量，本身这些变量不适合放在
         # inferstate的基类中，但是为了代码的简洁和方便，都放在基类中
         # 进行管理。注意这些成员变量只会在特定的模型和模式下才会生效。
@@ -90,7 +99,6 @@ def __init__(self):
         self.dp_input_split_sizes: List[List[int]] = None
 
     def init_some_extra_state(self, model):
-
         if self.is_prefill:
             (
                 self.b_q_seq_len,
@@ -103,7 +111,7 @@ def init_some_extra_state(self, model):
                 b_ready_cache_len=self.b_ready_cache_len,
                 b_seq_len=self.b_seq_len,
             )
-            self.b_start_loc = self.b1_cu_q_seq_len[0:-1]
+            self.b_q_start_loc = self.b1_cu_q_seq_len[0:-1]
         else:
             (
                 self.b_q_seq_len,
@@ -112,9 +120,17 @@ def init_some_extra_state(self, model):
                 self.b1_cu_kv_seq_len,
                 self.position_ids,
             ) = gen_decode_params(self.b_seq_len)
-            # TODO: check the correctness
-            self.max_kv_seq_len = self.max_len_in_batch
-            self.b_start_loc = self.b1_cu_kv_seq_len[0:-1]
+            self.b_kv_start_loc = self.b1_cu_kv_seq_len[0:-1]
+
+    def init_att_state(self):
+        if self.is_prefill:
+            self.prefill_att_state.init_state()
+            if self.prefill_att_state1 is not None:
+                self.prefill_att_state1.init_state()
+        else:
+            self.decode_att_state.init_state()
+            if self.decode_att_state1 is not None:
+                self.decode_att_state1.init_state()
 
     def copy_for_cuda_graph(self, new_infer_state: "InferStateInfo"):
         for attr_name, attr_value in vars(new_infer_state).items():
@@ -122,6 +138,10 @@ def copy_for_cuda_graph(self, new_infer_state: "InferStateInfo"):
                 attr_ = getattr(self, attr_name, None)
                 if attr_ is not None and attr_.data_ptr() != attr_value.data_ptr():
                     attr_.copy_(attr_value, non_blocking=True)
+
+        self.decode_att_state.copy_for_decode_cuda_graph(new_infer_state.decode_att_state)
+        if self.decode_att_state1 is not None:
+            self.decode_att_state1.copy_for_decode_cuda_graph(new_infer_state.decode_att_state1)
         return
 
     def prefill_dp_balance(self, input_ids: torch.Tensor):
diff --git a/lightllm/common/basemodel/layer_infer/post_layer_infer.py b/lightllm/common/basemodel/layer_infer/post_layer_infer.py
index d254eb510a..c7bae26ead 100644
--- a/lightllm/common/basemodel/layer_infer/post_layer_infer.py
+++ b/lightllm/common/basemodel/layer_infer/post_layer_infer.py
@@ -4,8 +4,7 @@
 class PostLayerInfer(BaseLayerInfer):
     """ """
 
-    def __init__(self, network_config, mode):
+    def __init__(self, network_config):
         super().__init__()
         self.network_config_ = network_config
-        self.mode = mode
         return
diff --git a/lightllm/common/basemodel/layer_infer/pre_layer_infer.py b/lightllm/common/basemodel/layer_infer/pre_layer_infer.py
index 3626346f20..e83fe89490 100644
--- a/lightllm/common/basemodel/layer_infer/pre_layer_infer.py
+++ b/lightllm/common/basemodel/layer_infer/pre_layer_infer.py
@@ -4,8 +4,7 @@
 class PreLayerInfer(BaseLayerInfer):
     """ """
 
-    def __init__(self, network_config, mode):
+    def __init__(self, network_config):
         super().__init__()
         self.network_config_ = network_config
-        self.mode = mode
         return
diff --git a/lightllm/common/basemodel/layer_infer/template/post_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/post_layer_infer_template.py
index fa7e96a694..1b7813fca3 100644
--- a/lightllm/common/basemodel/layer_infer/template/post_layer_infer_template.py
+++ b/lightllm/common/basemodel/layer_infer/template/post_layer_infer_template.py
@@ -6,8 +6,8 @@
 class PostLayerInferTpl(PostLayerInfer):
     """ """
 
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         self.eps_ = 1e-5
         self.vocab_size_ = network_config["vocab_size"]
         self.embed_dim_ = network_config["n_embed"]
diff --git a/lightllm/common/basemodel/layer_infer/template/pre_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/pre_layer_infer_template.py
index e7a0840794..04f8cda16b 100644
--- a/lightllm/common/basemodel/layer_infer/template/pre_layer_infer_template.py
+++ b/lightllm/common/basemodel/layer_infer/template/pre_layer_infer_template.py
@@ -5,8 +5,8 @@
 class PreLayerInferTpl(PreLayerInfer):
     """ """
 
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         self.eps_ = 1e-5
         return
 
diff --git a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_cohere_template.py b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_cohere_template.py
deleted file mode 100755
index 27f71a17ec..0000000000
--- a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_cohere_template.py
+++ /dev/null
@@ -1,136 +0,0 @@
-from functools import partial
-from typing import Tuple
-
-import torch
-import torch.distributed as dist
-
-from lightllm.common.basemodel.layer_infer.template.transformer_layer_infer_template import TransformerLayerInferTpl
-from lightllm.utils.infer_utils import mark_cost_time
-
-from ...infer_struct import InferStateInfo
-from ..transformer_layer_infer import TransformerLayerInfer
-from lightllm.distributed.communication_op import all_reduce
-
-
-class TransformerLayerCohereInferTpl(TransformerLayerInferTpl):
-    """ """
-
-    def __init__(self, layer_num, network_config, mode):
-        super().__init__(layer_num, network_config, mode)
-
-        self.use_qk_norm_ = self.network_config_.get("use_qk_norm", False)
-        return
-
-    def _att_norm(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
-        raise Exception("need to impl")
-
-    def _q_norm(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
-        raise Exception("need to impl")
-
-    def _k_norm(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
-        raise Exception("need to impl")
-
-    def _bind_norm(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
-        self._att_norm = partial(TransformerLayerCohereInferTpl._q_norm, self)
-        self._q_norm = partial(TransformerLayerCohereInferTpl._k_norm, self)
-        self._k_norm = partial(TransformerLayerCohereInferTpl._att_norm, self)
-
-    def _rotary_emb_fwd(self, q, kv, position_cos, position_sin):
-        raise Exception("need to impl")
-
-    def _bind_rotary_emb_fwd(self):
-        raise Exception("need to impl")
-
-    def _get_qkv(
-        self, input, infer_state: InferStateInfo, layer_weight
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        q = layer_weight.q_proj.mm(input.view(-1, self.embed_dim_))
-        cache_kv = layer_weight.kv_proj.mm(input.view(-1, self.embed_dim_)).view(
-            -1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_
-        )
-
-        if self.use_qk_norm_:
-            q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
-            k = cache_kv[:, 0 : self.tp_k_head_num_, :]
-            q = self._q_norm(q, infer_state, layer_weight)
-            cache_kv[:, 0 : self.tp_k_head_num_, :] = self._k_norm(k, infer_state, layer_weight)
-        self._rotary_emb_fwd(
-            q.view(-1, self.tp_q_head_num_, self.head_dim_),
-            cache_kv[:, 0 : self.tp_k_head_num_, :],
-            infer_state.position_cos,
-            infer_state.position_sin,
-        )
-        return q, cache_kv
-
-    def _context_attention_kernel(self, q, kv, infer_state: InferStateInfo, layer_weight, out=None) -> torch.Tensor:
-        raise Exception("need to impl")
-
-    def _token_attention_kernel(self, q, infer_state: InferStateInfo, layer_weight, out=None) -> torch.Tensor:
-        raise Exception("need to impl")
-
-    def _get_o(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
-        raise Exception("need to impl")
-
-    def _ffn(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
-        raise Exception("need to impl")
-
-    def _context_attention(self, input_embding, infer_state: InferStateInfo, layer_weight):
-        q, cache_kv = self._get_qkv(input_embding, infer_state, layer_weight)
-        self._post_cache_kv(cache_kv, infer_state, layer_weight)
-        o = self._context_attention_kernel(q, cache_kv, infer_state, layer_weight)
-        q = None
-        o = self._get_o(o, infer_state, layer_weight)
-        if self.tp_world_size_ > 1:
-            all_reduce(o, group=infer_state.dist_group, op=dist.ReduceOp.SUM, async_op=False)
-        infer_state._attn_out = o
-        return
-
-    def _context_ffn(self, input_embdings, infer_state: InferStateInfo, layer_weight):
-        ffn_out = self._ffn(input_embdings, infer_state, layer_weight)
-        if self.tp_world_size_ > 1:
-            all_reduce(ffn_out, group=infer_state.dist_group, op=dist.ReduceOp.SUM, async_op=False)
-        infer_state._ffn_out = ffn_out
-        return
-
-    def _token_attention(self, input_embding, infer_state: InferStateInfo, layer_weight):
-        q, cache_kv = self._get_qkv(input_embding, infer_state, layer_weight)
-        self._post_cache_kv(cache_kv, infer_state, layer_weight)
-        o = self._token_attention_kernel(q, infer_state, layer_weight)
-        q = None
-        o = self._get_o(o, infer_state, layer_weight)
-        if self.tp_world_size_ > 1:
-            all_reduce(o, group=infer_state.dist_group, op=dist.ReduceOp.SUM, async_op=False)
-        infer_state._attn_out = o
-        return
-
-    def _token_ffn(self, input_embdings, infer_state: InferStateInfo, layer_weight):
-        ffn_out = self._ffn(input_embdings, infer_state, layer_weight)
-        if self.tp_world_size_ > 1:
-            all_reduce(ffn_out, group=infer_state.dist_group, op=dist.ReduceOp.SUM, async_op=False)
-        infer_state._ffn_out = ffn_out
-        return
-
-    def _cohere_residual(self, input_embdings, infer_state: InferStateInfo):
-        # emb_addr = input_embdings.data_ptr()
-        # attn_out_addr = infer_state._attn_out.data_ptr()
-        # ffn_addr = infer_state._ffn_out.data_ptr()
-        # assert emb_addr != attn_out_addr
-        # assert emb_addr != ffn_addr
-        # assert attn_out_addr != ffn_addr
-        input_embdings.add_(
-            infer_state._attn_out.view(-1, self.embed_dim_) + infer_state._ffn_out.view(-1, self.embed_dim_)
-        )
-
-    def context_forward(self, input_embdings, infer_state: InferStateInfo, layer_weight):
-        input1 = self._att_norm(input_embdings, infer_state, layer_weight)
-        self._context_attention(input1, infer_state, layer_weight=layer_weight)
-        self._context_ffn(input1, infer_state, layer_weight)
-        self._cohere_residual(input_embdings, infer_state)
-        return input_embdings
-
-    def token_forward(self, input_embdings, infer_state: InferStateInfo, layer_weight):
-        input1 = self._att_norm(input_embdings, infer_state, layer_weight)
-        self._token_attention(input1, infer_state, layer_weight=layer_weight)
-        self._token_ffn(input1, infer_state, layer_weight)
-        self._cohere_residual(input_embdings, infer_state)
-        return input_embdings
diff --git a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py
index 436ca77d84..9153349c5d 100755
--- a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py
+++ b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py
@@ -3,8 +3,6 @@
 import torch.distributed as dist
 from ..transformer_layer_infer import TransformerLayerInfer
 from ...infer_struct import InferStateInfo
-from lightllm.utils.infer_utils import mark_cost_time
-from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv
 from lightllm.distributed import all_reduce
 from typing import Tuple
 from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor
@@ -13,8 +11,8 @@
 class TransformerLayerInferTpl(TransformerLayerInfer):
     """ """
 
-    def __init__(self, layer_num, network_config, mode):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         # need to set by subclass
         self.eps_ = 1e-5
         self.tp_q_head_num_ = -1
@@ -39,11 +37,11 @@ def _tpsp_get_qkv(self, input, infer_state: InferStateInfo, layer_weight) -> Tup
 
     def _post_cache_kv(self, cache_kv, infer_state: InferStateInfo, layer_weight):
         mem_manager = infer_state.mem_manager
-        self._copy_kv_to_mem_cache(cache_kv, infer_state.mem_index, mem_manager)
-        return
-
-    def _copy_kv_to_mem_cache(self, buffer, mem_index, mem_manager):
-        destindex_copy_kv(buffer, mem_index, mem_manager.kv_buffer[self.layer_num_])
+        mem_manager.copy_kv_to_mem_manager(
+            layer_index=self.layer_num_,
+            mem_index=infer_state.mem_index,
+            kv=cache_kv,
+        )
         return
 
     def _context_attention_kernel(self, q, kv, infer_state: InferStateInfo, layer_weight, out=None) -> torch.Tensor:
diff --git a/lightllm/common/basemodel/layer_infer/transformer_layer_infer.py b/lightllm/common/basemodel/layer_infer/transformer_layer_infer.py
index 7350531bbb..53daffcddf 100644
--- a/lightllm/common/basemodel/layer_infer/transformer_layer_infer.py
+++ b/lightllm/common/basemodel/layer_infer/transformer_layer_infer.py
@@ -4,9 +4,8 @@
 class TransformerLayerInfer(BaseLayerInfer):
     """ """
 
-    def __init__(self, layer_num, network_config, mode):
+    def __init__(self, layer_num, network_config):
         super().__init__()
         self.layer_num_ = layer_num
         self.network_config_ = network_config
-        self.mode = mode
         return
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/__init__.py b/lightllm/common/basemodel/layer_weights/meta_weights/__init__.py
index 0fa02780cb..8e884012d5 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/__init__.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/__init__.py
@@ -1,13 +1,12 @@
 from .base_weight import BaseWeight
 from .mm_weight import (
-    MMWeightPack,
     MMWeightTpl,
     ROWMMWeight,
-    COLMMWeight,
+    KVROWNMMWeight,
     ROWBMMWeight,
+    COLMMWeight,
 )
-from .norm_weight import NoTpGEMMANormWeight, TpVitPadNormWeight, NoTpNormWeight, TpHeadNormWeight
-from .fused_moe_weight_tp import create_tp_moe_wegiht_obj
-from .fused_moe_weight_ep import FusedMoeWeightEP
+from .norm_weight import TpRMSNormWeight, RMSNormWeight, LayerNormWeight, NoTpGEMMANormWeight, QKRMSNORMWeight
 from .embedding_weight import EmbeddingWeight, LMHeadWeight, NoTpPosEmbeddingWeight
 from .att_sink_weight import TpAttSinkWeight
+from .fused_moe.fused_moe_weight import FusedMoeWeight
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/att_sink_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/att_sink_weight.py
index 3f8e1f50ab..32d59e66e7 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/att_sink_weight.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/att_sink_weight.py
@@ -10,6 +10,7 @@ def __init__(self, weight_name: str, data_type):
         self.weight_name = weight_name
         self.data_type_ = data_type
         self.weight: torch.Tensor = None
+        # TODO: add create weight function
 
     def load_hf_weights(self, weights: Dict[str, torch.Tensor]):
         if self.weight_name not in weights or self.weight is not None:
@@ -18,6 +19,3 @@ def load_hf_weights(self, weights: Dict[str, torch.Tensor]):
         t_weight = weights[self.weight_name]
         start_head_index, end_head_index = self._get_head_tp_split_params(weight=t_weight)
         self.weight = t_weight[start_head_index:end_head_index].to(self.data_type_).cuda(get_current_device_id())
-
-    def verify_load(self):
-        return self.weight is not None
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py
index 2cd8ea6ae0..da03887862 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py
@@ -6,6 +6,7 @@
 
 class BaseWeight(ABC):
     def __init__(self):
+        super().__init__()
         pass
 
     @abstractmethod
@@ -13,12 +14,17 @@ def load_hf_weights(self, weights):
         pass
 
     @abstractmethod
-    def verify_load(self) -> bool:
+    def _create_weight(self):
+        pass
+
+    @abstractmethod
+    def verify_load(self):
         pass
 
 
 class BaseWeightTpl(BaseWeight):
     def __init__(self, tp_rank: int = None, tp_world_size: int = None, data_type: torch.dtype = None):
+        super().__init__()
         self.tp_world_size_ = tp_world_size if tp_world_size is not None else get_dp_world_size()
         self.tp_rank_ = tp_rank if tp_rank is not None else get_current_rank_in_dp()
         self.device_id_ = get_current_device_id()
@@ -27,32 +33,8 @@ def __init__(self, tp_rank: int = None, tp_world_size: int = None, data_type: to
     def load_hf_weights(self, weights):
         raise NotImplementedError("load_hf_weights must implement this method")
 
-    def verify_load(self) -> bool:
+    def verify_load(self):
         raise NotImplementedError("verify_load must implement this method")
 
-    def _get_head_tp_split_params(self, weight: torch.Tensor) -> Tuple[int, int]:
-        """
-        Docstring for _get_head_tp_split_params,
-        一个常用的tp 划分head获取head_index 范围的功能函数, 一些继承类可能会使用。
-        :param self: Description
-        :param weight: Description
-        :type weight: torch.Tensor
-        :return: Description
-        :rtype: Tuple[int, int]
-        """
-        assert weight.ndim == 2
-
-        all_head_num = weight.shape[0]
-        tp_head_num = all_head_num // self.tp_world_size_
-
-        if tp_head_num > 0:
-            start_head_index = self.tp_rank_ * tp_head_num
-            end_head_index = (self.tp_rank_ + 1) * tp_head_num
-        else:
-            # 当 tp_world_size 大于 all_head_num 时的特殊处理
-            scale_size = self.tp_world_size_ // all_head_num
-            assert self.tp_world_size_ % all_head_num == 0
-            start_head_index = self.tp_rank_ // scale_size
-            end_head_index = start_head_index + 1
-
-        return start_head_index, end_head_index
+    def _create_weight(self) -> bool:
+        raise NotImplementedError("create_weight must implement this method")
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/embedding_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/embedding_weight.py
index fc018267fa..d4e03d0a1f 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/embedding_weight.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/embedding_weight.py
@@ -2,46 +2,64 @@
 import numpy as np
 from typing import Dict, Optional
 from .base_weight import BaseWeightTpl
-from lightllm.utils.dist_utils import get_current_device_id
+from .platform_op import PlatformAwareOp
 from lightllm.common.basemodel.triton_kernel.embedding import embedding as embedding_kernel
-from lightllm.utils.log_utils import init_logger
+from lightllm.utils.dist_utils import get_dp_world_size, get_current_rank_in_dp
 
-logger = init_logger(__name__)
 
-
-class EmbeddingWeight(BaseWeightTpl):
-    def __init__(self, weight_name, data_type):
+class EmbeddingWeight(BaseWeightTpl, PlatformAwareOp):
+    def __init__(self, dim: int, vocab_size: int, weight_name: str, data_type: torch.dtype):
         super().__init__()
+        self.dim = dim
+        self.vocab_size = vocab_size
+        self.tp_world_size_ = get_dp_world_size()
+        self.tp_rank_ = get_current_rank_in_dp()
+        # 计算 split_indexes
+        split_indexes = np.linspace(0, self.vocab_size, self.tp_world_size_ + 1, dtype=np.int64)
+        self.tp_vocab_start_id = int(split_indexes[self.tp_rank_])
+        self.tp_vocab_end_id = int(split_indexes[self.tp_rank_ + 1])
         self.weight_name: str = weight_name
         self.data_type_ = data_type
-        self.weight: torch.Tensor = None
+        self._create_weight()
+
+    def _create_weight(self):
+        tp_vocab_size = self.tp_vocab_end_id - self.tp_vocab_start_id
+        self.weight: torch.Tensor = torch.empty(tp_vocab_size, self.dim, dtype=self.data_type_, device=self.device_id_)
+        self.weight.load_ok = False
 
     def load_hf_weights(self, weights: Dict[str, torch.Tensor]):
-        if self.weight_name not in weights or self.weight is not None:
+        if self.weight_name not in weights:
             return
-
         t_weight = weights[self.weight_name]
         # init some params
-        self.vocab_size = len(t_weight)
-        split_indexes = np.linspace(0, self.vocab_size, self.tp_world_size_ + 1, dtype=np.int64)
-        self.tp_vocab_start_id = int(split_indexes[self.tp_rank_])
-        self.tp_vocab_end_id = int(split_indexes[self.tp_rank_ + 1])
-
-        logger.info(f"loaded weight vocab_size: {self.vocab_size}")
-
-        self.weight = (
-            t_weight[self.tp_vocab_start_id : self.tp_vocab_end_id, :].to(self.data_type_).cuda(get_current_device_id())
-        )
+        loaded_vocab_size = len(t_weight)
+        assert (
+            loaded_vocab_size == self.vocab_size
+        ), f"loaded weight vocab_size: {loaded_vocab_size} != expected vocab_size: {self.vocab_size}"
+        self.weight.copy_(t_weight[self.tp_vocab_start_id : self.tp_vocab_end_id, :].to(self.data_type_))
+        self.weight.load_ok = True
 
     def verify_load(self):
-        return self.weight is not None
-
-    def embedding(self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty):
+        return self.weight.load_ok
+
+    def _native_forward(
+        self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None, _alloc_func=torch.empty
+    ) -> torch.Tensor:
+        adjusted_ids = input_ids - self.tp_vocab_start_id
+        adjusted_ids = torch.clamp(adjusted_ids, 0, self.weight.shape[0] - 1)
+        result = torch.nn.functional.embedding(adjusted_ids, self.weight)
+        if out is not None:
+            out.copy_(result)
+            return out
+        return result
+
+    def _triton_forward(
+        self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
         if out is None:
             out = alloc_func(
                 (input_ids.shape[0], self.weight.shape[1]), dtype=self.weight.dtype, device=self.weight.device
             )
-
         embedding_kernel(
             input_ids=input_ids,
             weight=self.weight,
@@ -49,10 +67,86 @@ def embedding(self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None,
             vob_end_id=self.tp_vocab_end_id,
             out=out,
         )
-
         return out
 
-    def lm_head(self, input: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty):
+    def _cuda_forward(
+        self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        return self._triton_forward(input_ids=input_ids, out=out, alloc_func=alloc_func)
+
+    def _musa_forward(
+        self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        # triton implementation is supported by musa.
+        return self._triton_forward(input_ids=input_ids, out=out, alloc_func=alloc_func)
+
+    def __call__(
+        self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        return self._forward(input_ids=input_ids, out=out, alloc_func=alloc_func)
+
+
+class LMHeadWeight(BaseWeightTpl, PlatformAwareOp):
+    def __init__(
+        self,
+        dim: int,
+        vocab_size: int,
+        weight_name: str,
+        data_type: torch.dtype,
+        embedding_weight: Optional[EmbeddingWeight] = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.vocab_size = vocab_size
+        self.tp_world_size_ = get_dp_world_size()
+        self.tp_rank_ = get_current_rank_in_dp()
+        # 计算 split_indexes
+        split_indexes = np.linspace(0, self.vocab_size, self.tp_world_size_ + 1, dtype=np.int64)
+        self.tp_vocab_start_id = int(split_indexes[self.tp_rank_])
+        self.tp_vocab_end_id = int(split_indexes[self.tp_rank_ + 1])
+        self.weight_name: str = weight_name
+        self.data_type_ = data_type
+        self._embedding_weight = embedding_weight
+        self._create_weight()
+
+    def _create_weight(self):
+        if self._embedding_weight is not None:
+            self.weight = self._embedding_weight.weight
+            return
+        tp_vocab_size = self.tp_vocab_end_id - self.tp_vocab_start_id
+        self.weight: torch.Tensor = torch.empty(tp_vocab_size, self.dim, dtype=self.data_type_, device=self.device_id_)
+        self.weight.load_ok = False
+
+    def load_hf_weights(self, weights: Dict[str, torch.Tensor]):
+        # When set tile_embedding=True, no need to load - EmbeddingWeight already loaded it
+        if self._embedding_weight is not None:
+            return
+        if self.weight_name not in weights:
+            return
+        t_weight = weights[self.weight_name]
+        loaded_vocab_size = len(t_weight)
+        assert (
+            loaded_vocab_size == self.vocab_size
+        ), f"loaded weight vocab_size: {loaded_vocab_size} != expected vocab_size: {self.vocab_size}"
+        self.weight.copy_(t_weight[self.tp_vocab_start_id : self.tp_vocab_end_id, :].to(self.data_type_))
+        self.weight.load_ok = True
+
+    def verify_load(self):
+        return self.weight.load_ok
+
+    def _native_forward(
+        self, input: torch.Tensor, out: Optional[torch.Tensor] = None, _alloc_func=torch.empty
+    ) -> torch.Tensor:
+        assert input.ndim == 2
+        result = torch.mm(self.weight, input)
+        if out is not None:
+            out.copy_(result)
+            return out
+        return result
+
+    def _cuda_forward(
+        self, input: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
         assert input.ndim == 2
         if out is None:
             out = alloc_func(
@@ -60,49 +154,71 @@ def lm_head(self, input: torch.Tensor, out: Optional[torch.Tensor] = None, alloc
                 dtype=input.dtype,
                 device=input.device,
             )
-
         torch.mm(self.weight, input, out=out)
         return out
 
-
-class LMHeadWeight(EmbeddingWeight):
-    def __init__(self, weight_name, data_type):
-        super().__init__(weight_name, data_type)
+    def __call__(self, input: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty) -> torch.Tensor:
+        return self._forward(input=input, out=out, alloc_func=alloc_func)
 
 
-class NoTpPosEmbeddingWeight(BaseWeightTpl):
-    def __init__(self, weight_name, data_type):
+class NoTpPosEmbeddingWeight(BaseWeightTpl, PlatformAwareOp):
+    def __init__(self, dim: int, max_position_embeddings: int, weight_name: str, data_type: torch.dtype):
         super().__init__()
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
         self.weight_name: str = weight_name
         self.data_type_ = data_type
-        self.weight: torch.Tensor = None
         self.tp_world_size_ = 1
         self.tp_rank_ = 0
+        self._create_weight()
+
+    def _create_weight(self):
+        self.weight: torch.Tensor = torch.empty(
+            self.max_position_embeddings, self.dim, dtype=self.data_type_, device=self.device_id_
+        )
+        self.weight.load_ok = False
 
     def load_hf_weights(self, weights: Dict[str, torch.Tensor]):
-        if self.weight_name not in weights or self.weight is not None:
+        if self.weight_name not in weights:
             return
-
         t_weight = weights[self.weight_name]
-        self.weight = t_weight.to(self.data_type_).cuda(get_current_device_id())
-        self.end_position_id: int = t_weight.shape[0]
-        logger.info(f"loaded weight end_position_id: {self.end_position_id}")
+        loaded_max_position_embeddings = t_weight.shape[0]
+        assert (
+            loaded_max_position_embeddings == self.max_position_embeddings
+        ), f"max_position_embeddings: {loaded_max_position_embeddings} != expected: {self.max_position_embeddings}"
+        self.weight.copy_(t_weight.to(self.data_type_))
+        self.weight.load_ok = True
 
     def verify_load(self):
-        return self.weight is not None
-
-    def embedding(self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty):
+        return self.weight.load_ok
+
+    def _native_forward(
+        self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None, _alloc_func=torch.empty
+    ) -> torch.Tensor:
+        # Use PyTorch native embedding
+        result = torch.nn.functional.embedding(input_ids, self.weight)
+        if out is not None:
+            out.copy_(result)
+            return out
+        return result
+
+    def _cuda_forward(
+        self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
         if out is None:
             out = alloc_func(
                 (input_ids.shape[0], self.weight.shape[1]), dtype=self.weight.dtype, device=self.weight.device
             )
-
         embedding_kernel(
             input_ids=input_ids,
             weight=self.weight,
             vob_start_id=0,
-            vob_end_id=self.end_position_id,
+            vob_end_id=self.max_position_embeddings,
             out=out,
         )
-
         return out
+
+    def __call__(
+        self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        return self._forward(input_ids=input_ids, out=out, alloc_func=alloc_func)
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_ep_redundancy.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_redundancy.py
similarity index 95%
rename from lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_ep_redundancy.py
rename to lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_redundancy.py
index b53200d4c8..a31cd18803 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_ep_redundancy.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_redundancy.py
@@ -1,6 +1,6 @@
 import numpy as np
 import torch
-from .fused_moe_weight_ep import FusedMoeWeightEP
+from .fused_moe_weight import FusedMoeWeight
 from lightllm.utils.log_utils import init_logger
 from typing import Dict
 
@@ -10,7 +10,7 @@
 class FusedMoeWeightEPAutoRedundancy:
     def __init__(
         self,
-        ep_fused_moe_weight: FusedMoeWeightEP,
+        ep_fused_moe_weight: FusedMoeWeight,
     ) -> None:
         super().__init__()
         self._ep_w = ep_fused_moe_weight
@@ -102,12 +102,12 @@ def _fuse(self):
                 inter_shape, hidden_size = self.w2_list[0].shape[0], self.w2_list[0].shape[1]
                 w2 = torch._utils._flatten_dense_tensors(self.w2_list).view(len(self.w2_list), inter_shape, hidden_size)
                 if not self._ep_w.quantized_weight and self._ep_w.quant_method is not None:
-                    qw1, qw1_scale, qw1_zero_point = self._ep_w.quant_method.quantize(w1)
-                    qw2, qw2_scale, qw2_zero_point = self._ep_w.quant_method.quantize(w2)
-                    self.w1[0] = qw1
-                    self.w1[1] = qw1_scale
-                    self.w2[0] = qw2
-                    self.w2[1] = qw2_scale
+                    qw1_pack = self._ep_w.quant_method.quantize(w1)
+                    qw2_pack = self._ep_w.quant_method.quantize(w2)
+                    self.w1[0] = qw1_pack.weight
+                    self.w1[1] = qw1_pack.weight_scale
+                    self.w2[0] = qw2_pack.weight
+                    self.w2[1] = qw2_pack.weight_scale
                 else:
                     self.w1[0] = w1
                     self.w2[0] = w2
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py
new file mode 100644
index 0000000000..6a1bd0ca45
--- /dev/null
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py
@@ -0,0 +1,363 @@
+import torch
+import threading
+from typing import Dict, Any, Optional, Tuple
+from lightllm.common.basemodel.layer_weights.meta_weights.base_weight import BaseWeightTpl
+from lightllm.common.quantization.quantize_method import WeightPack
+from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_slicer import (
+    get_row_slice_mixin,
+    get_col_slice_mixin,
+)
+from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import select_fuse_moe_impl
+from lightllm.common.quantization.quantize_method import QuantizationMethod
+from lightllm.utils.envs_utils import get_redundancy_expert_ids, get_redundancy_expert_num, get_env_start_args
+from lightllm.utils.dist_utils import get_global_world_size, get_global_rank
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class FusedMoeWeight(BaseWeightTpl):
+    def __init__(
+        self,
+        gate_proj_name: str,
+        down_proj_name: str,
+        up_proj_name: str,
+        e_score_correction_bias_name: str,
+        weight_prefix: str,
+        n_routed_experts: int,
+        hidden_size: int,
+        moe_intermediate_size: int,
+        data_type: torch.dtype,
+        quant_method: QuantizationMethod = None,
+        num_fused_shared_experts: int = 0,
+        layer_num: int = 0,
+        network_config: Dict[str, Any] = None,
+    ) -> None:
+        super().__init__(data_type=data_type)
+        self.w1_weight_name = gate_proj_name
+        self.w2_weight_name = down_proj_name
+        self.w3_weight_name = up_proj_name
+        self.e_score_correction_bias_name = e_score_correction_bias_name
+        self.weight_prefix = weight_prefix
+        self.layer_num_ = layer_num
+        self.global_rank_ = get_global_rank()
+        self.global_world_size = get_global_world_size()
+        self.hidden_size = hidden_size
+        self.moe_intermediate_size = moe_intermediate_size
+        self.quant_method = quant_method
+        self.row_slicer = get_row_slice_mixin(
+            self.quant_method.method_name, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_
+        )
+        self.col_slicer = get_col_slice_mixin(
+            self.quant_method.method_name, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_
+        )
+        assert num_fused_shared_experts in [0, 1], "num_fused_shared_experts can only support 0 or 1 now."
+        self.enable_ep_moe = get_env_start_args().enable_ep_moe
+        self.n_routed_experts = n_routed_experts
+        self.num_fused_shared_experts = num_fused_shared_experts
+        self._init_config(network_config)
+        self._init_redundancy_expert_params()
+        self._init_parallel_params()
+        self.fuse_moe_impl = select_fuse_moe_impl(self.quant_method, self.enable_ep_moe)(
+            n_routed_experts=self.n_routed_experts,
+            num_fused_shared_experts=self.num_fused_shared_experts,
+            routed_scaling_factor=self.routed_scaling_factor,
+            quant_method=self.quant_method,
+            redundancy_expert_num=self.redundancy_expert_num,
+            redundancy_expert_ids_tensor=self.redundancy_expert_ids_tensor,
+            routed_expert_counter_tensor=self.routed_expert_counter_tensor,
+            auto_update_redundancy_expert=self.auto_update_redundancy_expert,
+        )
+        self.lock = threading.Lock()
+        self._create_weight()
+
+    def _init_config(self, network_config: Dict[str, Any]):
+        self.n_group = network_config.get("n_group", 0)
+        self.use_grouped_topk = self.n_group > 0
+        self.norm_topk_prob = network_config["norm_topk_prob"]
+        self.topk_group = network_config.get("topk_group", 0)
+        self.num_experts_per_tok = network_config["num_experts_per_tok"]
+        self.routed_scaling_factor = network_config.get("routed_scaling_factor", 1.0)
+        self.scoring_func = network_config.get("scoring_func", "softmax")
+
+    def _init_redundancy_expert_params(self):
+        self.redundancy_expert_num = get_redundancy_expert_num()
+        self.redundancy_expert_ids = get_redundancy_expert_ids(self.layer_num_)
+        self.auto_update_redundancy_expert: bool = get_env_start_args().auto_update_redundancy_expert
+        self.redundancy_expert_ids_tensor = torch.tensor(self.redundancy_expert_ids, dtype=torch.int64, device="cuda")
+        self.routed_expert_counter_tensor = torch.zeros((self.n_routed_experts,), dtype=torch.int64, device="cuda")
+        # TODO: find out the reason of failure of deepep when redundancy_expert_num is 1.
+        assert self.redundancy_expert_num != 1, "redundancy_expert_num can not be 1 for some unknown hang of deepep."
+
+    def _init_parallel_params(self):
+        self.local_n_routed_experts = self.n_routed_experts + self.num_fused_shared_experts
+        self.split_inter_size = self.moe_intermediate_size // self.tp_world_size_
+        if self.enable_ep_moe:
+            assert self.num_fused_shared_experts == 0, "num_fused_shared_experts must be 0 when enable_ep_moe"
+            logger.info(
+                f"global_rank {self.global_rank_} layerindex {self.layer_num_} "
+                f"redundancy_expertids: {self.redundancy_expert_ids}"
+            )
+            self.local_n_routed_experts = self.n_routed_experts // self.global_world_size + self.redundancy_expert_num
+            self.split_inter_size = self.moe_intermediate_size
+            n_experts_per_rank = self.n_routed_experts // self.global_world_size
+            start_expert_id = self.global_rank_ * n_experts_per_rank
+            self.local_expert_ids = (
+                list(range(start_expert_id, start_expert_id + n_experts_per_rank)) + self.redundancy_expert_ids
+            )
+            self.expert_idx_to_local_idx = {
+                expert_idx: expert_idx - start_expert_id for expert_idx in self.local_expert_ids[:n_experts_per_rank]
+            }
+            self.redundancy_expert_idx_to_local_idx = {
+                redundancy_expert_idx: n_experts_per_rank + i
+                for (i, redundancy_expert_idx) in enumerate(self.redundancy_expert_ids)
+            }
+        else:
+            self.local_expert_ids = list(range(self.n_routed_experts + self.num_fused_shared_experts))
+            self.expert_idx_to_local_idx = {expert_idx: i for (i, expert_idx) in enumerate(self.local_expert_ids)}
+            self.rexpert_idx_to_local_idx = {}
+
+    def experts(
+        self,
+        input_tensor: torch.Tensor,
+        router_logits: torch.Tensor,
+        top_k: int,
+        renormalize: bool,
+        use_grouped_topk: bool,
+        topk_group: int,
+        num_expert_group: int,
+        is_prefill: Optional[bool] = None,
+    ):
+        """Backward compatible method that routes to platform-specific implementation."""
+        return self.fuse_moe_impl(
+            input_tensor=input_tensor,
+            router_logits=router_logits,
+            w13=self.w13,
+            w2=self.w2,
+            correction_bias=self.e_score_correction_bias,
+            scoring_func=self.scoring_func,
+            top_k=top_k,
+            renormalize=renormalize,
+            use_grouped_topk=use_grouped_topk,
+            topk_group=topk_group,
+            num_expert_group=num_expert_group,
+            is_prefill=is_prefill,
+        )
+
+    def low_latency_dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+    ):
+        assert self.enable_ep_moe, "low_latency_dispatch is only supported when enable_ep_moe is True"
+        return self.fuse_moe_impl.low_latency_dispatch(
+            hidden_states=hidden_states,
+            router_logits=router_logits,
+            e_score_correction_bias=self.e_score_correction_bias,
+            use_grouped_topk=self.use_grouped_topk,
+            num_experts_per_tok=self.num_experts_per_tok,
+            norm_topk_prob=self.norm_topk_prob,
+            topk_group=self.topk_group,
+            n_group=self.n_group,
+            scoring_func=self.scoring_func,
+        )
+
+    def select_experts_and_quant_input(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+    ):
+        assert self.enable_ep_moe, "select_experts_and_quant_input is only supported when enable_ep_moe is True"
+        return self.fuse_moe_impl.select_experts_and_quant_input(
+            hidden_states=hidden_states,
+            router_logits=router_logits,
+            e_score_correction_bias=self.e_score_correction_bias,
+            w13=self.w13,
+            use_grouped_topk=self.use_grouped_topk,
+            num_experts_per_tok=self.num_experts_per_tok,
+            norm_topk_prob=self.norm_topk_prob,
+            topk_group=self.topk_group,
+            n_group=self.n_group,
+            scoring_func=self.scoring_func,
+        )
+
+    def dispatch(
+        self,
+        qinput_tensor: Tuple[torch.Tensor],
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+        overlap_event: Optional[Any] = None,
+    ):
+        assert self.enable_ep_moe, "dispatch is only supported when enable_ep_moe is True"
+        return self.fuse_moe_impl.dispatch(
+            qinput_tensor=qinput_tensor,
+            topk_idx=topk_idx,
+            topk_weights=topk_weights,
+            overlap_event=overlap_event,
+        )
+
+    def masked_group_gemm(
+        self, recv_x: Tuple[torch.Tensor], masked_m: torch.Tensor, dtype: torch.dtype, expected_m: int
+    ):
+        assert self.enable_ep_moe, "masked_group_gemm is only supported when enable_ep_moe is True"
+        return self.fuse_moe_impl.masked_group_gemm(
+            recv_x=recv_x,
+            w13=self.w13,
+            w2=self.w2,
+            masked_m=masked_m,
+            dtype=dtype,
+            expected_m=expected_m,
+        )
+
+    def prefilled_group_gemm(
+        self,
+        num_recv_tokens_per_expert_list,
+        recv_x: Tuple[torch.Tensor],
+        recv_topk_idx: torch.Tensor,
+        recv_topk_weights: torch.Tensor,
+        hidden_dtype=torch.bfloat16,
+    ):
+        assert self.enable_ep_moe, "prefilled_group_gemm is only supported when enable_ep_moe is True"
+        return self.fuse_moe_impl.prefilled_group_gemm(
+            num_recv_tokens_per_expert_list=num_recv_tokens_per_expert_list,
+            recv_x=recv_x,
+            recv_topk_idx=recv_topk_idx,
+            recv_topk_weights=recv_topk_weights,
+            w13=self.w13,
+            w2=self.w2,
+            hidden_dtype=hidden_dtype,
+        )
+
+    def low_latency_combine(
+        self,
+        gemm_out_b: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+        handle: Any,
+    ):
+        assert self.enable_ep_moe, "low_latency_combine is only supported when enable_ep_moe is True"
+        return self.fuse_moe_impl.low_latency_combine(
+            gemm_out_b=gemm_out_b,
+            topk_idx=topk_idx,
+            topk_weights=topk_weights,
+            handle=handle,
+        )
+
+    def combine(
+        self,
+        gemm_out_b: torch.Tensor,
+        handle: Any,
+        overlap_event: Optional[Any] = None,
+    ):
+        assert self.enable_ep_moe, "combine is only supported when enable_ep_moe is True"
+        return self.fuse_moe_impl.combine(
+            gemm_out_b=gemm_out_b,
+            handle=handle,
+            overlap_event=overlap_event,
+        )
+
+    def load_hf_weights(self, weights):
+        # Load bias
+        if self.e_score_correction_bias_name in weights:
+            self.e_score_correction_bias.copy_(weights[self.e_score_correction_bias_name])
+        self._load_weight(self.expert_idx_to_local_idx, weights)
+        if self.redundancy_expert_num > 0:
+            self._load_weight(self.redundancy_expert_idx_to_local_idx, weights)
+
+    def verify_load(self):
+        return True
+
+    def _create_weight(self):
+        intermediate_size = self.split_inter_size
+        self.e_score_correction_bias = None
+        # Create e_score_correction_bias
+        if self.e_score_correction_bias_name:
+            self.e_score_correction_bias = torch.empty(
+                (self.n_routed_experts,),
+                dtype=self.data_type_,
+                device=f"cuda:{self.device_id_}",
+            )
+
+        self.w13: WeightPack = self.quant_method.create_weight(
+            out_dim=intermediate_size * 2,
+            in_dim=self.hidden_size,
+            dtype=self.data_type_,
+            device_id=self.device_id_,
+            num_experts=self.local_n_routed_experts,
+        )
+        self.w2: WeightPack = self.quant_method.create_weight(
+            out_dim=self.hidden_size,
+            in_dim=intermediate_size,
+            dtype=self.data_type_,
+            device_id=self.device_id_,
+            num_experts=self.local_n_routed_experts,
+        )
+        self.load_cnt = 0
+
+    def _load_weight(self, expert_idx_to_local_idx: Dict[int, int], weights: Dict[str, torch.Tensor]):
+
+        # Load each expert with TP slicing
+        for expert_idx, local_expert_idx in expert_idx_to_local_idx.items():
+            with self.lock:
+                self._load_expert(
+                    expert_idx, local_expert_idx, weights, type="weight", suffix=self.quant_method.weight_suffix
+                )
+            if self.w13.weight_scale is not None:
+                with self.lock:
+                    self._load_expert(
+                        expert_idx,
+                        local_expert_idx,
+                        weights,
+                        type="weight_scale",
+                        suffix=self.quant_method.weight_scale_suffix,
+                    )
+            if self.w13.weight_zero_point is not None:
+                with self.lock:
+                    self._load_expert(
+                        expert_idx,
+                        local_expert_idx,
+                        weights,
+                        type="weight_zero_point",
+                        suffix=self.quant_method.weight_zero_point_suffix,
+                    )
+
+    def _load_expert(
+        self,
+        expert_idx: int,
+        local_expert_idx: int,
+        weights: Dict[str, torch.Tensor],
+        type: str,
+        suffix: str = "weight",
+    ):
+        w1_weight = f"{self.weight_prefix}.{expert_idx}.{self.w1_weight_name}.{suffix}"
+        w2_weight = f"{self.weight_prefix}.{expert_idx}.{self.w2_weight_name}.{suffix}"
+        w3_weight = f"{self.weight_prefix}.{expert_idx}.{self.w3_weight_name}.{suffix}"
+        intermediate_size = self.split_inter_size
+        load_func, slice_func = self._get_load_and_slice_func(type, is_row=True)
+        if w1_weight in weights:
+            load_func(slice_func(weights[w1_weight]), self.w13.get_expert(local_expert_idx), start_idx=0)
+        if w3_weight in weights:
+            load_func(
+                slice_func(weights[w3_weight]), self.w13.get_expert(local_expert_idx), start_idx=intermediate_size
+            )
+        load_func, slice_func = self._get_load_and_slice_func(type, is_row=False)
+        if w2_weight in weights:
+            load_func(slice_func(weights[w2_weight]), self.w2.get_expert(local_expert_idx), start_idx=0)
+
+    def _load_weight_func(self, weight: torch.Tensor, weight_pack: WeightPack, start_idx: int = 0):
+        if self.quant_method.weight_need_quanted(weight):
+            self.quant_method.quantize(weight, weight_pack, start_idx)
+        else:
+            self.quant_method.load_weight(weight, weight_pack, start_idx)
+
+    def _get_load_and_slice_func(self, type: str, is_row: bool = True):
+        if is_row:
+            slicer = self.row_slicer
+        else:
+            slicer = self.col_slicer
+        if type == "weight":
+            return self._load_weight_func, slicer._slice_weight
+        elif type == "weight_scale":
+            return getattr(self.quant_method, "load_weight_scale"), slicer._slice_weight_scale
+        elif type == "weight_zero_point":
+            return getattr(self.quant_method, "load_weight_zero_point"), slicer._slice_weight_zero_point
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/gpt_oss_fused_moe_weight_tp.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py
similarity index 98%
rename from lightllm/common/basemodel/layer_weights/meta_weights/gpt_oss_fused_moe_weight_tp.py
rename to lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py
index df72cc6208..e7748b1dfd 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/gpt_oss_fused_moe_weight_tp.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py
@@ -3,7 +3,7 @@
 import threading
 from typing import Optional, Tuple, List, Dict, Any
 
-from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe_weight_tp import FusedMoeWeightTP
+from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight
 from lightllm.utils.dist_utils import get_current_rank_in_dp, get_current_device_id
 from lightllm.common.quantization import Quantcfg
 from lightllm.utils.log_utils import init_logger
@@ -30,7 +30,7 @@
 ]
 
 
-class GPTOSSFusedMoeWeightTP(FusedMoeWeightTP):
+class GPTOSSFusedMoeWeightTP(FusedMoeWeight):
     def __init__(
         self,
         gate_up_proj_name: str,  # diff with FusedMoeWeightTP
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py
new file mode 100644
index 0000000000..67bb90e4ef
--- /dev/null
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py
@@ -0,0 +1,14 @@
+from lightllm.common.quantization.quantize_method import QuantizationMethod
+from .triton_impl import FuseMoeTriton
+from .marlin_impl import FuseMoeMarlin
+from .deepgemm_impl import FuseMoeDeepGEMM
+
+
+def select_fuse_moe_impl(quant_method: QuantizationMethod, enable_ep_moe: bool):
+    if enable_ep_moe:
+        return FuseMoeDeepGEMM
+
+    if quant_method.method_name == "awq_marlin":
+        return FuseMoeMarlin
+    else:
+        return FuseMoeTriton
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py
new file mode 100644
index 0000000000..c56cd4da31
--- /dev/null
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py
@@ -0,0 +1,66 @@
+import torch
+from abc import abstractmethod
+from lightllm.common.quantization.quantize_method import (
+    WeightPack,
+    QuantizationMethod,
+)
+from typing import Optional
+from lightllm.utils.dist_utils import (
+    get_global_rank,
+    get_global_world_size,
+)
+
+
+class FuseMoeBaseImpl:
+    def __init__(
+        self,
+        n_routed_experts: int,
+        num_fused_shared_experts: int,
+        routed_scaling_factor: float,
+        quant_method: QuantizationMethod,
+        redundancy_expert_num: int,
+        redundancy_expert_ids_tensor: torch.Tensor,
+        routed_expert_counter_tensor: torch.Tensor,
+        auto_update_redundancy_expert: bool,
+    ):
+        self.n_routed_experts = n_routed_experts
+        self.num_fused_shared_experts = num_fused_shared_experts
+        self.routed_scaling_factor = routed_scaling_factor
+        self.quant_method = quant_method
+        self.global_rank_ = get_global_rank()
+        self.global_world_size_ = get_global_world_size()
+        self.ep_n_routed_experts = self.n_routed_experts // self.global_world_size_
+        self.total_expert_num_contain_redundancy = (
+            self.n_routed_experts + redundancy_expert_num * self.global_world_size_
+        )
+
+        # redundancy expert related
+        self.redundancy_expert_num = redundancy_expert_num
+        self.redundancy_expert_ids_tensor = redundancy_expert_ids_tensor
+        self.routed_expert_counter_tensor = routed_expert_counter_tensor
+        self.auto_update_redundancy_expert = auto_update_redundancy_expert
+
+        # workspace for kernel optimization
+        self.workspace = self.create_workspace()
+
+    @abstractmethod
+    def create_workspace(self):
+        pass
+
+    @abstractmethod
+    def __call__(
+        self,
+        input_tensor: torch.Tensor,
+        router_logits: torch.Tensor,
+        w13: WeightPack,
+        w2: WeightPack,
+        correction_bias: Optional[torch.Tensor],
+        scoring_func: str,
+        top_k: int,
+        renormalize: bool,
+        use_grouped_topk: bool,
+        topk_group: int,
+        num_expert_group: int,
+        is_prefill: Optional[bool] = None,
+    ):
+        pass
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py
new file mode 100644
index 0000000000..f00d572d9d
--- /dev/null
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py
@@ -0,0 +1,336 @@
+import torch
+from typing import Optional, Tuple, Any
+from .triton_impl import FuseMoeTriton
+from lightllm.distributed import dist_group_manager
+from lightllm.common.triton_utils.autotuner import Autotuner
+from lightllm.common.quantization.quantize_method import WeightPack
+from lightllm.utils.envs_utils import get_deepep_num_max_dispatch_tokens_per_rank
+from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import (
+    fused_experts_impl,
+    masked_group_gemm,
+    _deepgemm_grouped_fp8_nt_contiguous,
+)
+from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import (
+    per_token_group_quant_fp8,
+    tma_align_input_scale,
+)
+from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_scatter_gather import ep_scatter, ep_gather
+from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
+from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair
+
+
+class FuseMoeDeepGEMM(FuseMoeTriton):
+    def _select_experts(
+        self,
+        input_tensor: torch.Tensor,
+        router_logits: torch.Tensor,
+        correction_bias: Optional[torch.Tensor],
+        top_k: int,
+        renormalize: bool,
+        use_grouped_topk: bool,
+        topk_group: int,
+        num_expert_group: int,
+        scoring_func: str,
+    ):
+        """Select experts and return topk weights and ids."""
+        from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import select_experts
+
+        topk_weights, topk_ids = select_experts(
+            hidden_states=input_tensor,
+            router_logits=router_logits,
+            correction_bias=correction_bias,
+            use_grouped_topk=use_grouped_topk,
+            top_k=top_k,
+            renormalize=renormalize,
+            topk_group=topk_group,
+            num_expert_group=num_expert_group,
+            scoring_func=scoring_func,
+        )
+        topk_weights.mul_(self.routed_scaling_factor)
+        if self.redundancy_expert_num > 0:
+            redundancy_topk_ids_repair(
+                topk_ids=topk_ids,
+                redundancy_expert_ids=self.redundancy_expert_ids_tensor,
+                ep_expert_num=self.ep_n_routed_experts,
+                global_rank=self.global_rank_,
+                expert_counter=self.routed_expert_counter_tensor,
+                enable_counter=self.auto_update_redundancy_expert,
+            )
+        return topk_weights, topk_ids
+
+    def _fused_experts(
+        self,
+        input_tensor: torch.Tensor,
+        w13: WeightPack,
+        w2: WeightPack,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        router_logits: Optional[torch.Tensor] = None,
+        is_prefill: Optional[bool] = None,
+    ):
+
+        w13_weight, w13_scale = w13.weight, w13.weight_scale
+        w2_weight, w2_scale = w2.weight, w2.weight_scale
+        use_fp8_w8a8 = self.quant_method.method_name != "none"
+        output = fused_experts_impl(
+            hidden_states=input_tensor,
+            w1=w13_weight,
+            w2=w2_weight,
+            topk_weights=topk_weights,
+            topk_idx=topk_ids.to(torch.long),
+            num_experts=self.total_expert_num_contain_redundancy,  # number of all experts contain redundancy
+            buffer=dist_group_manager.ep_buffer,
+            is_prefill=is_prefill,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_fp8_all2all=use_fp8_w8a8,
+            use_int8_w8a16=False,  # default to False
+            w1_scale=w13_scale,
+            w2_scale=w2_scale,
+            previous_event=None,  # for overlap
+        )
+        return output
+
+    def low_latency_dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        e_score_correction_bias: torch.Tensor,
+        use_grouped_topk: bool,
+        num_experts_per_tok: int,
+        norm_topk_prob: bool,
+        topk_group: int,
+        n_group: int,
+        scoring_func: str,
+    ):
+        topk_weights, topk_idx = self._select_experts(
+            input_tensor=hidden_states,
+            router_logits=router_logits,
+            correction_bias=e_score_correction_bias,
+            use_grouped_topk=use_grouped_topk,
+            top_k=num_experts_per_tok,
+            renormalize=norm_topk_prob,
+            topk_group=topk_group,
+            num_expert_group=n_group,
+            scoring_func=scoring_func,
+        )
+
+        topk_idx = topk_idx.to(torch.long)
+        num_max_dispatch_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank()
+        use_fp8_w8a8 = self.quant_method.method_name != "none"
+        recv_x, masked_m, handle, event, hook = dist_group_manager.ep_buffer.low_latency_dispatch(
+            hidden_states,
+            topk_idx,
+            num_max_dispatch_tokens_per_rank,
+            self.total_expert_num_contain_redundancy,
+            use_fp8=use_fp8_w8a8,
+            async_finish=False,
+            return_recv_hook=True,
+        )
+        return recv_x, masked_m, topk_idx, topk_weights, handle, hook
+
+    def select_experts_and_quant_input(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        e_score_correction_bias: torch.Tensor,
+        w13: WeightPack,
+        use_grouped_topk: bool,
+        num_experts_per_tok: int,
+        norm_topk_prob: bool,
+        topk_group: int,
+        n_group: int,
+        scoring_func: str,
+    ):
+        topk_weights, topk_idx = self._select_experts(
+            input_tensor=hidden_states,
+            router_logits=router_logits,
+            correction_bias=e_score_correction_bias,
+            use_grouped_topk=use_grouped_topk,
+            top_k=num_experts_per_tok,
+            renormalize=norm_topk_prob,
+            topk_group=topk_group,
+            num_expert_group=n_group,
+            scoring_func=scoring_func,
+        )
+        w13_weight, w13_scale = w13.weight, w13.weight_scale
+        block_size_k = 0
+        if w13_weight.ndim == 3:
+            block_size_k = w13_weight.shape[2] // w13_scale.shape[2]
+        assert block_size_k == 128, "block_size_k must be 128"
+        qinput_tensor, input_scale = per_token_group_quant_fp8(hidden_states, block_size_k, dtype=w13_weight.dtype)
+        return topk_weights, topk_idx.to(torch.long), (qinput_tensor, input_scale)
+
+    def dispatch(
+        self,
+        qinput_tensor: Tuple[torch.Tensor],
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+        overlap_event: Optional[Any] = None,
+    ):
+        buffer = dist_group_manager.ep_buffer
+        # get_dispatch_layout
+        (
+            num_tokens_per_rank,
+            num_tokens_per_rdma_rank,
+            num_tokens_per_expert,
+            is_token_in_rank,
+            previous_event,
+        ) = buffer.get_dispatch_layout(
+            topk_idx,
+            self.total_expert_num_contain_redundancy,
+            previous_event=overlap_event,
+            async_finish=True,
+            allocate_on_comm_stream=True,
+        )
+        recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event = buffer.dispatch(
+            qinput_tensor,
+            topk_idx=topk_idx,
+            topk_weights=topk_weights,
+            num_tokens_per_rank=num_tokens_per_rank,
+            num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
+            is_token_in_rank=is_token_in_rank,
+            num_tokens_per_expert=num_tokens_per_expert,
+            previous_event=previous_event,
+            async_finish=True,
+            allocate_on_comm_stream=True,
+            expert_alignment=128,
+        )
+
+        def hook():
+            event.current_stream_wait()
+
+        return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, hook
+
+    def masked_group_gemm(
+        self,
+        recv_x: Tuple[torch.Tensor],
+        w13: WeightPack,
+        w2: WeightPack,
+        masked_m: torch.Tensor,
+        dtype: torch.dtype,
+        expected_m: int,
+    ):
+        w13_weight, w13_scale = w13.weight, w13.weight_scale
+        w2_weight, w2_scale = w2.weight, w2.weight_scale
+        return masked_group_gemm(
+            recv_x, masked_m, dtype, w13_weight, w13_scale, w2_weight, w2_scale, expected_m=expected_m
+        )
+
+    def prefilled_group_gemm(
+        self,
+        num_recv_tokens_per_expert_list,
+        recv_x: Tuple[torch.Tensor],
+        recv_topk_idx: torch.Tensor,
+        recv_topk_weights: torch.Tensor,
+        w13: WeightPack,
+        w2: WeightPack,
+        hidden_dtype=torch.bfloat16,
+    ):
+        device = recv_x[0].device
+        w13_weight, w13_scale = w13.weight, w13.weight_scale
+        w2_weight, w2_scale = w2.weight, w2.weight_scale
+        _, K = recv_x[0].shape
+        _, N, _ = w13_weight.shape
+        block_size = self.quant_method.block_size
+        # scatter
+        all_tokens = sum(num_recv_tokens_per_expert_list)  # calcu padding all nums.
+        # gather_out shape [recive_num_tokens, hidden]
+        gather_out = torch.empty_like(recv_x[0], device=device, dtype=hidden_dtype)
+        if all_tokens > 0:
+            input_tensor = [
+                torch.empty((all_tokens, K), device=device, dtype=recv_x[0].dtype),
+                torch.empty((all_tokens, K // 128), device=device, dtype=torch.float32),
+            ]
+            # when m_indices is filled ok.
+            # m_indices show token use which expert, example, [0, 0, 0, 0, .... 1, 1, 1, 1,...., cur_expert_num - 1, ..]
+            # the count of 0 is num_recv_tokens_per_expert_list[0], the count of 1 is num_recv_tokens_per_expert_list[1]
+            # ...
+            m_indices = torch.empty(all_tokens, device=device, dtype=torch.int32)
+            # output_index shape [recive_num_tokens, topk_num]
+            # output_index use to show the token index in input_tensor
+            output_index = torch.empty_like(recv_topk_idx)
+
+            num_recv_tokens_per_expert = torch.tensor(
+                num_recv_tokens_per_expert_list, dtype=torch.int32, pin_memory=True, device="cpu"
+            ).cuda(non_blocking=True)
+
+            expert_start_loc = torch.empty_like(num_recv_tokens_per_expert)
+
+            ep_scatter(
+                recv_x[0],
+                recv_x[1],
+                recv_topk_idx,
+                num_recv_tokens_per_expert,
+                expert_start_loc,
+                input_tensor[0],
+                input_tensor[1],
+                m_indices,
+                output_index,
+            )
+            input_tensor[1] = tma_align_input_scale(input_tensor[1])
+            # groupgemm (contiguous layout)
+            gemm_out_a = torch.empty((all_tokens, N), device=device, dtype=hidden_dtype)
+
+            _deepgemm_grouped_fp8_nt_contiguous(input_tensor, (w13_weight, w13_scale), gemm_out_a, m_indices)
+
+            # silu_and_mul_fwd + qaunt
+            # TODO fused kernel
+            silu_out = torch.empty((all_tokens, N // 2), device=device, dtype=hidden_dtype)
+
+            silu_and_mul_fwd(gemm_out_a.view(-1, N), silu_out)
+            qsilu_out, qsilu_out_scale = per_token_group_quant_fp8(
+                silu_out, block_size, dtype=w13_weight.dtype, column_major_scales=True, scale_tma_aligned=True
+            )
+
+            # groupgemm (contiguous layout)
+            gemm_out_b = torch.empty((all_tokens, K), device=device, dtype=hidden_dtype)
+
+            _deepgemm_grouped_fp8_nt_contiguous(
+                (qsilu_out, qsilu_out_scale), (w2_weight, w2_scale), gemm_out_b, m_indices
+            )
+            # gather and local reduce
+            ep_gather(gemm_out_b, recv_topk_idx, recv_topk_weights, output_index, gather_out)
+        else:
+            ######################################## warning ##################################################
+            # here is used to match autotune feature, make moe model run same triton kernel in different rank.
+            # in some special case, one rank will recv 0 token, so add a token to make it run triton kernel.
+            if Autotuner.is_autotune_warmup():
+                _gemm_out_a = torch.zeros((1, N), device=device, dtype=hidden_dtype)
+                _silu_out = torch.zeros((1, N // 2), device=device, dtype=hidden_dtype)
+                silu_and_mul_fwd(_gemm_out_a.view(-1, N), _silu_out)
+                _gemm_out_a, _silu_out = None, None
+
+        return gather_out
+
+    def low_latency_combine(
+        self,
+        gemm_out_b: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+        handle: Any,
+    ):
+        combined_x, event_overlap, hook = dist_group_manager.ep_buffer.low_latency_combine(
+            gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=True
+        )
+        return combined_x, hook
+
+    def combine(
+        self,
+        gemm_out_b: torch.Tensor,
+        handle: Any,
+        overlap_event: Optional[Any] = None,
+    ):
+        # normal combine
+        combined_x, _, event = dist_group_manager.ep_buffer.combine(
+            gemm_out_b,
+            handle,
+            topk_weights=None,
+            async_finish=True,
+            previous_event=overlap_event,
+            allocate_on_comm_stream=True,
+        )
+
+        def hook():
+            event.current_stream_wait()
+
+        return combined_x, hook
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py
new file mode 100644
index 0000000000..bdccbc0eef
--- /dev/null
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py
@@ -0,0 +1,56 @@
+import torch
+from .triton_impl import FuseMoeTriton
+from lightllm.common.quantization.quantize_method import (
+    WeightPack,
+)
+from typing import Optional
+
+
+class FuseMoeMarlin(FuseMoeTriton):
+    def create_workspace(self):
+        from lightllm.utils.vllm_utils import HAS_VLLM
+
+        assert HAS_VLLM, "moe awq marlin quantization requires kernels of vllm"
+        from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+            marlin_make_workspace_new,
+        )
+
+        return marlin_make_workspace_new(torch.device("cuda"), 4)
+
+    def _fused_experts(
+        self,
+        input_tensor: torch.Tensor,
+        w13: WeightPack,
+        w2: WeightPack,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        router_logits: Optional[torch.Tensor] = None,
+        is_prefill: Optional[bool] = None,
+    ):
+
+        w1_weight, w1_scale, w1_zero_point = w13.weight, w13.weight_scale, w13.weight_zero_point
+        w2_weight, w2_scale, w2_zero_point = w2.weight, w2.weight_scale, w2.weight_zero_point
+
+        from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
+
+        fused_marlin_moe(
+            input_tensor,
+            w1_weight,
+            w2_weight,
+            None,
+            None,
+            w1_scale,
+            w2_scale,
+            router_logits,
+            topk_weights,
+            topk_ids,
+            quant_type_id=self.quant_method.vllm_quant_type.id,
+            apply_router_weight_on_input=False,
+            global_num_experts=-1,
+            expert_map=None,
+            w1_zeros=w1_zero_point,
+            w2_zeros=w2_zero_point,
+            workspace=self.workspace,
+            inplace=True,
+        )
+        return input_tensor
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py
new file mode 100644
index 0000000000..8bcdb4bf90
--- /dev/null
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py
@@ -0,0 +1,148 @@
+import torch
+from typing import Optional
+from lightllm.common.quantization.no_quant import WeightPack
+from lightllm.common.quantization.quantize_method import QuantizationMethod
+from .base_impl import FuseMoeBaseImpl
+
+
+class FuseMoeTriton(FuseMoeBaseImpl):
+    def __init__(
+        self,
+        n_routed_experts: int,
+        num_fused_shared_experts: int,
+        routed_scaling_factor: float,
+        quant_method: QuantizationMethod,
+        redundancy_expert_num: int,
+        redundancy_expert_ids_tensor: torch.Tensor,
+        routed_expert_counter_tensor: torch.Tensor,
+        auto_update_redundancy_expert: bool,
+    ):
+        super().__init__(
+            n_routed_experts=n_routed_experts,
+            num_fused_shared_experts=num_fused_shared_experts,
+            routed_scaling_factor=routed_scaling_factor,
+            quant_method=quant_method,
+            redundancy_expert_num=redundancy_expert_num,
+            redundancy_expert_ids_tensor=redundancy_expert_ids_tensor,
+            routed_expert_counter_tensor=routed_expert_counter_tensor,
+            auto_update_redundancy_expert=auto_update_redundancy_expert,
+        )
+
+    def create_workspace(self):
+        return None
+
+    def _select_experts(
+        self,
+        input_tensor: torch.Tensor,
+        router_logits: torch.Tensor,
+        correction_bias: Optional[torch.Tensor],
+        top_k: int,
+        renormalize: bool,
+        use_grouped_topk: bool,
+        topk_group: int,
+        num_expert_group: int,
+        scoring_func: str,
+    ):
+        """Select experts and return topk weights and ids."""
+        from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import select_experts
+
+        topk_weights, topk_ids = select_experts(
+            hidden_states=input_tensor,
+            router_logits=router_logits,
+            correction_bias=correction_bias,
+            use_grouped_topk=use_grouped_topk,
+            top_k=top_k,
+            renormalize=renormalize,
+            topk_group=topk_group,
+            num_expert_group=num_expert_group,
+            scoring_func=scoring_func,
+        )
+        topk_weights.mul_(self.routed_scaling_factor)
+        if self.num_fused_shared_experts > 0:
+            pad_topk_ids = (
+                torch.arange(
+                    start=self.n_routed_experts,
+                    end=self.n_routed_experts + self.num_fused_shared_experts,
+                    step=1,
+                    dtype=topk_ids.dtype,
+                    device="cuda",
+                )
+                .view(1, self.num_fused_shared_experts)
+                .repeat(topk_ids.shape[0], 1)
+            )
+            pad_topk_weights = torch.full(
+                (topk_weights.shape[0], self.num_fused_shared_experts),
+                fill_value=1.0,
+                device="cuda",
+                dtype=topk_weights.dtype,
+            )
+
+            topk_ids = torch.cat([topk_ids, pad_topk_ids], dim=1)
+            topk_weights = torch.cat([topk_weights, pad_topk_weights], dim=1)
+        return topk_weights, topk_ids
+
+    def _fused_experts(
+        self,
+        input_tensor: torch.Tensor,
+        w13: WeightPack,
+        w2: WeightPack,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        router_logits: Optional[torch.Tensor] = None,
+        is_prefill: bool = False,
+    ):
+        w13_weight, w13_scale = w13.weight, w13.weight_scale
+        w2_weight, w2_scale = w2.weight, w2.weight_scale
+        use_fp8_w8a8 = w13_weight.dtype == torch.float8_e4m3fn
+
+        from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import fused_experts
+
+        fused_experts(
+            hidden_states=input_tensor,
+            w1=w13_weight,
+            w2=w2_weight,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            inplace=True,
+            use_fp8_w8a8=use_fp8_w8a8,
+            w1_scale=w13_scale,
+            w2_scale=w2_scale,
+        )
+        return input_tensor
+
+    def __call__(
+        self,
+        input_tensor: torch.Tensor,
+        router_logits: torch.Tensor,
+        w13: WeightPack,
+        w2: WeightPack,
+        correction_bias: Optional[torch.Tensor],
+        scoring_func: str,
+        top_k: int,
+        renormalize: bool,
+        use_grouped_topk: bool,
+        topk_group: int,
+        num_expert_group: int,
+        is_prefill: Optional[bool] = None,
+    ):
+        topk_weights, topk_ids = self._select_experts(
+            input_tensor=input_tensor,
+            router_logits=router_logits,
+            correction_bias=correction_bias,
+            top_k=top_k,
+            renormalize=renormalize,
+            use_grouped_topk=use_grouped_topk,
+            topk_group=topk_group,
+            num_expert_group=num_expert_group,
+            scoring_func=scoring_func,
+        )
+        output = self._fused_experts(
+            input_tensor=input_tensor,
+            w13=w13,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            router_logits=router_logits,
+            is_prefill=is_prefill,
+        )
+        return output
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_ep.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_ep.py
deleted file mode 100644
index 7dc5b5fdcc..0000000000
--- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_ep.py
+++ /dev/null
@@ -1,540 +0,0 @@
-import os
-import torch
-import threading
-from typing import Optional, Tuple, List, Dict, Any
-from lightllm.utils.dist_utils import get_global_world_size, get_global_rank, get_current_device_id
-from .base_weight import BaseWeight
-from lightllm.common.fused_moe.grouped_fused_moe_ep import (
-    fused_experts_impl,
-    masked_group_gemm,
-    _deepgemm_grouped_fp8_nt_contiguous,
-)
-from lightllm.common.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
-from lightllm.distributed import dist_group_manager
-from lightllm.common.fused_moe.topk_select import select_experts
-from lightllm.utils.envs_utils import get_deepep_num_max_dispatch_tokens_per_rank
-from lightllm.utils.envs_utils import get_redundancy_expert_ids, get_redundancy_expert_num
-from lightllm.utils.envs_utils import get_env_start_args
-from lightllm.common.quantization.triton_quant.fp8.fp8act_quant_kernel import (
-    per_token_group_quant_fp8,
-    tma_align_input_scale,
-)
-from lightllm.common.fused_moe.deepep_scatter_gather import ep_scatter, ep_gather
-from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair
-from lightllm.utils.log_utils import init_logger
-from lightllm.common.triton_utils.autotuner import Autotuner
-
-
-logger = init_logger(__name__)
-
-
-class FusedMoeWeightEP(BaseWeight):
-    def __init__(
-        self,
-        gate_proj_name: str,
-        down_proj_name: str,
-        up_proj_name: str,
-        e_score_correction_bias_name: str,
-        weight_prefix: str,
-        n_routed_experts: int,
-        data_type: torch.dtype,
-        network_config: Dict[str, Any],
-        layer_num: int,
-        quant_cfg=None,
-    ) -> None:
-        super().__init__()
-
-        self.layer_num = layer_num
-        self.quant_method = quant_cfg.get_quant_method(layer_num, "fused_moe")
-        self.quantized_weight = quant_cfg.quantized_weight
-        if self.quant_method is not None:
-            self.weight_scale_suffix = self.quant_method.weight_scale_suffix
-            self.quant_method.is_moe = True
-            block_size = 1
-            if hasattr(self.quant_method, "block_size"):
-                block_size = self.quant_method.block_size
-            self.block_size = block_size
-
-        self.weight_prefix = weight_prefix
-        self.w1_weight_name = gate_proj_name
-        self.w2_weight_name = down_proj_name
-        self.w3_weight_name = up_proj_name
-        self.e_score_correction_bias_name = e_score_correction_bias_name
-        self.n_routed_experts = n_routed_experts
-        self.data_type_ = data_type
-
-        global_world_size = get_global_world_size()
-        self.global_rank_ = get_global_rank()
-        self.redundancy_expert_num = get_redundancy_expert_num()
-        self.redundancy_expert_ids = get_redundancy_expert_ids(layer_num)
-        logger.info(
-            f"global_rank {self.global_rank_} layerindex {layer_num} redundancy_expertids: {self.redundancy_expert_ids}"
-        )
-        self.redundancy_expert_ids_tensor = torch.tensor(self.redundancy_expert_ids, dtype=torch.int64, device="cuda")
-        self.routed_expert_counter_tensor = torch.zeros((self.n_routed_experts,), dtype=torch.int64, device="cuda")
-        self.total_expert_num_contain_redundancy = (
-            self.n_routed_experts + self.redundancy_expert_num * global_world_size
-        )
-        assert self.n_routed_experts % global_world_size == 0
-        self.ep_n_routed_experts = self.n_routed_experts // global_world_size
-        ep_load_expert_num = self.ep_n_routed_experts + self.redundancy_expert_num
-        self.experts_up_projs = [None] * ep_load_expert_num
-        self.experts_gate_projs = [None] * ep_load_expert_num
-        self.experts_up_proj_scales = [None] * ep_load_expert_num
-        self.experts_gate_proj_scales = [None] * ep_load_expert_num
-        self.e_score_correction_bias = None
-        self.w2_list = [None] * ep_load_expert_num
-        self.w2_scale_list = [None] * ep_load_expert_num
-        self.scoring_func = network_config.get("scoring_func", "softmax")
-        self.w1 = [None, None]  # weight, weight_scale
-        self.w2 = [None, None]  # weight, weight_scale
-        self.use_fp8_w8a8 = self.quant_method is not None
-        network_config["n_group"] = network_config.get("n_group", 0)
-        self.num_experts_per_tok = network_config["num_experts_per_tok"]
-        self.use_grouped_topk = network_config["n_group"] > 0
-        self.norm_topk_prob = network_config["norm_topk_prob"]
-        self.n_group = network_config["n_group"]
-        network_config["topk_group"] = network_config.get("topk_group", 0)
-        self.topk_group = network_config["topk_group"]
-        network_config["routed_scaling_factor"] = network_config.get("routed_scaling_factor", 1.0)
-        self.routed_scaling_factor = network_config["routed_scaling_factor"]
-
-        self.lock = threading.Lock()
-        # init buffer
-
-        # auto update redundancy expert vars
-        self.auto_update_redundancy_expert: bool = get_env_start_args().auto_update_redundancy_expert
-
-    def experts(
-        self,
-        input_tensor,
-        router_logits,
-        top_k,
-        renormalize,
-        use_grouped_topk,
-        topk_group,
-        num_expert_group,
-        is_prefill,
-    ):
-        topk_weights, topk_ids = select_experts(
-            hidden_states=input_tensor,
-            router_logits=router_logits,
-            correction_bias=self.e_score_correction_bias,
-            use_grouped_topk=use_grouped_topk,
-            top_k=top_k,
-            renormalize=renormalize,
-            topk_group=topk_group,
-            num_expert_group=num_expert_group,
-            scoring_func=self.scoring_func,
-        )
-        topk_weights.mul_(self.routed_scaling_factor)
-
-        if self.redundancy_expert_num > 0:
-            redundancy_topk_ids_repair(
-                topk_ids=topk_ids,
-                redundancy_expert_ids=self.redundancy_expert_ids_tensor,
-                ep_expert_num=self.ep_n_routed_experts,
-                global_rank=self.global_rank_,
-                expert_counter=self.routed_expert_counter_tensor,
-                enable_counter=self.auto_update_redundancy_expert,
-            )
-
-        w1, w1_scale = self.w1
-        w2, w2_scale = self.w2
-        return fused_experts_impl(
-            hidden_states=input_tensor,
-            w1=w1,
-            w2=w2,
-            topk_weights=topk_weights,
-            topk_idx=topk_ids.to(torch.long),
-            num_experts=self.total_expert_num_contain_redundancy,  # number of all experts contain redundancy
-            buffer=dist_group_manager.ep_buffer,
-            is_prefill=is_prefill,
-            use_fp8_w8a8=self.use_fp8_w8a8,
-            use_fp8_all2all=self.use_fp8_w8a8,
-            use_int8_w8a16=False,  # default to False
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            previous_event=None,  # for overlap
-        )
-
-    def low_latency_dispatch(
-        self,
-        hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
-    ):
-
-        topk_weights, topk_idx = select_experts(
-            hidden_states=hidden_states,
-            router_logits=router_logits,
-            correction_bias=self.e_score_correction_bias,
-            use_grouped_topk=self.use_grouped_topk,
-            top_k=self.num_experts_per_tok,
-            renormalize=self.norm_topk_prob,
-            topk_group=self.topk_group,
-            num_expert_group=self.n_group,
-            scoring_func=self.scoring_func,
-        )
-        topk_weights.mul_(self.routed_scaling_factor)
-
-        if self.redundancy_expert_num > 0:
-            redundancy_topk_ids_repair(
-                topk_ids=topk_idx,
-                redundancy_expert_ids=self.redundancy_expert_ids_tensor,
-                ep_expert_num=self.ep_n_routed_experts,
-                global_rank=self.global_rank_,
-                expert_counter=self.routed_expert_counter_tensor,
-                enable_counter=self.auto_update_redundancy_expert,
-            )
-
-        topk_idx = topk_idx.to(torch.long)
-        num_max_dispatch_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank()
-        recv_x, masked_m, handle, event, hook = dist_group_manager.ep_buffer.low_latency_dispatch(
-            hidden_states,
-            topk_idx,
-            num_max_dispatch_tokens_per_rank,
-            self.total_expert_num_contain_redundancy,
-            use_fp8=self.use_fp8_w8a8,
-            async_finish=False,
-            return_recv_hook=True,
-        )
-        return recv_x, masked_m, topk_idx, topk_weights, handle, hook
-
-    def select_experts_and_quant_input(
-        self,
-        hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
-    ):
-        topk_weights, topk_idx = select_experts(
-            hidden_states=hidden_states,
-            router_logits=router_logits,
-            correction_bias=self.e_score_correction_bias,
-            use_grouped_topk=self.use_grouped_topk,
-            top_k=self.num_experts_per_tok,
-            renormalize=self.norm_topk_prob,
-            topk_group=self.topk_group,
-            num_expert_group=self.n_group,
-            scoring_func=self.scoring_func,
-        )
-        topk_weights.mul_(self.routed_scaling_factor)
-        if self.redundancy_expert_num > 0:
-            redundancy_topk_ids_repair(
-                topk_ids=topk_idx,
-                redundancy_expert_ids=self.redundancy_expert_ids_tensor,
-                ep_expert_num=self.ep_n_routed_experts,
-                global_rank=self.global_rank_,
-                expert_counter=self.routed_expert_counter_tensor,
-                enable_counter=self.auto_update_redundancy_expert,
-            )
-        M, K = hidden_states.shape
-        w1, w1_scale = self.w1
-        block_size_k = 0
-        if w1.ndim == 3:
-            block_size_k = w1.shape[2] // w1_scale.shape[2]
-        assert block_size_k == 128, "block_size_k must be 128"
-        qinput_tensor, input_scale = per_token_group_quant_fp8(hidden_states, block_size_k, dtype=w1.dtype)
-        return topk_weights, topk_idx.to(torch.long), (qinput_tensor, input_scale)
-
-    def dispatch(
-        self,
-        qinput_tensor: Tuple[torch.Tensor],
-        topk_idx: torch.Tensor,
-        topk_weights: torch.Tensor,
-        overlap_event: Optional[Any] = None,
-    ):
-        buffer = dist_group_manager.ep_buffer
-        # get_dispatch_layout
-        (
-            num_tokens_per_rank,
-            num_tokens_per_rdma_rank,
-            num_tokens_per_expert,
-            is_token_in_rank,
-            previous_event,
-        ) = buffer.get_dispatch_layout(
-            topk_idx,
-            self.total_expert_num_contain_redundancy,
-            previous_event=overlap_event,
-            async_finish=True,
-            allocate_on_comm_stream=True,
-        )
-        recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event = buffer.dispatch(
-            qinput_tensor,
-            topk_idx=topk_idx,
-            topk_weights=topk_weights,
-            num_tokens_per_rank=num_tokens_per_rank,
-            num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
-            is_token_in_rank=is_token_in_rank,
-            num_tokens_per_expert=num_tokens_per_expert,
-            previous_event=previous_event,
-            async_finish=True,
-            allocate_on_comm_stream=True,
-            expert_alignment=128,
-        )
-
-        def hook():
-            event.current_stream_wait()
-
-        return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, hook
-
-    def masked_group_gemm(
-        self, recv_x: Tuple[torch.Tensor], masked_m: torch.Tensor, dtype: torch.dtype, expected_m: int
-    ):
-        w1, w1_scale = self.w1
-        w2, w2_scale = self.w2
-        return masked_group_gemm(recv_x, masked_m, dtype, w1, w1_scale, w2, w2_scale, expected_m=expected_m)
-
-    def prefilled_group_gemm(
-        self,
-        num_recv_tokens_per_expert_list,
-        recv_x: Tuple[torch.Tensor],
-        recv_topk_idx: torch.Tensor,
-        recv_topk_weights: torch.Tensor,
-        hidden_dtype=torch.bfloat16,
-    ):
-        device = recv_x[0].device
-        w1, w1_scale = self.w1
-        w2, w2_scale = self.w2
-        _, K = recv_x[0].shape
-        _, N, _ = w1.shape
-        # scatter
-        all_tokens = sum(num_recv_tokens_per_expert_list)  # calcu padding all nums.
-        # gather_out shape [recive_num_tokens, hidden]
-        gather_out = torch.empty_like(recv_x[0], device=device, dtype=hidden_dtype)
-        if all_tokens > 0:
-            input_tensor = [
-                torch.empty((all_tokens, K), device=device, dtype=recv_x[0].dtype),
-                torch.empty((all_tokens, K // 128), device=device, dtype=torch.float32),
-            ]
-            # when m_indices is filled ok.
-            # m_indices show token use which expert, example, [0, 0, 0, 0, .... 1, 1, 1, 1,...., cur_expert_num - 1, ..]
-            # the count of 0 is num_recv_tokens_per_expert_list[0], the count of 1 is num_recv_tokens_per_expert_list[1]
-            # ...
-            m_indices = torch.empty(all_tokens, device=device, dtype=torch.int32)
-            # output_index shape [recive_num_tokens, topk_num]
-            # output_index use to show the token index in input_tensor
-            output_index = torch.empty_like(recv_topk_idx)
-
-            num_recv_tokens_per_expert = torch.tensor(
-                num_recv_tokens_per_expert_list, dtype=torch.int32, pin_memory=True, device="cpu"
-            ).cuda(non_blocking=True)
-
-            expert_start_loc = torch.empty_like(num_recv_tokens_per_expert)
-
-            ep_scatter(
-                recv_x[0],
-                recv_x[1],
-                recv_topk_idx,
-                num_recv_tokens_per_expert,
-                expert_start_loc,
-                input_tensor[0],
-                input_tensor[1],
-                m_indices,
-                output_index,
-            )
-            input_tensor[1] = tma_align_input_scale(input_tensor[1])
-            # groupgemm (contiguous layout)
-            gemm_out_a = torch.empty((all_tokens, N), device=device, dtype=hidden_dtype)
-
-            _deepgemm_grouped_fp8_nt_contiguous(input_tensor, (w1, w1_scale), gemm_out_a, m_indices)
-
-            # silu_and_mul_fwd + qaunt
-            # TODO fused kernel
-            silu_out = torch.empty((all_tokens, N // 2), device=device, dtype=hidden_dtype)
-
-            silu_and_mul_fwd(gemm_out_a.view(-1, N), silu_out)
-            qsilu_out, qsilu_out_scale = per_token_group_quant_fp8(
-                silu_out, self.block_size, dtype=w1.dtype, column_major_scales=True, scale_tma_aligned=True
-            )
-
-            # groupgemm (contiguous layout)
-            gemm_out_b = torch.empty((all_tokens, K), device=device, dtype=hidden_dtype)
-
-            _deepgemm_grouped_fp8_nt_contiguous((qsilu_out, qsilu_out_scale), (w2, w2_scale), gemm_out_b, m_indices)
-            # gather and local reduce
-            ep_gather(gemm_out_b, recv_topk_idx, recv_topk_weights, output_index, gather_out)
-        else:
-            ######################################## warning ##################################################
-            # here is used to match autotune feature, make moe model run same triton kernel in different rank.
-            # in some special case, one rank will recv 0 token, so add a token to make it run triton kernel.
-            if Autotuner.is_autotune_warmup():
-                _gemm_out_a = torch.zeros((1, N), device=device, dtype=hidden_dtype)
-                _silu_out = torch.zeros((1, N // 2), device=device, dtype=hidden_dtype)
-                silu_and_mul_fwd(_gemm_out_a.view(-1, N), _silu_out)
-                _gemm_out_a, _silu_out = None, None
-
-        return gather_out
-
-    def low_latency_combine(
-        self,
-        gemm_out_b: torch.Tensor,
-        topk_idx: torch.Tensor,
-        topk_weights: torch.Tensor,
-        handle: Any,
-    ):
-        combined_x, event_overlap, hook = dist_group_manager.ep_buffer.low_latency_combine(
-            gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=True
-        )
-        return combined_x, hook
-
-    def combine(
-        self,
-        gemm_out_b: torch.Tensor,
-        handle: Any,
-        overlap_event: Optional[Any] = None,
-    ):
-        # normal combine
-        combined_x, _, event = dist_group_manager.ep_buffer.combine(
-            gemm_out_b,
-            handle,
-            topk_weights=None,
-            async_finish=True,
-            previous_event=overlap_event,
-            allocate_on_comm_stream=True,
-        )
-
-        def hook():
-            event.current_stream_wait()
-
-        return combined_x, hook
-
-    def _fuse(self):
-        if self.quantized_weight:
-            self._fuse_weight_scale()
-        with self.lock:
-            if (
-                hasattr(self, "experts_up_projs")
-                and None not in self.experts_up_projs
-                and None not in self.experts_gate_projs
-                and None not in self.w2_list
-            ):
-                gate_out_dim, gate_in_dim = self.experts_gate_projs[0].shape
-                up_out_dim, up_in_dim = self.experts_up_projs[0].shape
-                assert gate_in_dim == up_in_dim
-                dtype = self.experts_gate_projs[0].dtype
-                total_expert_num = self.ep_n_routed_experts + self.redundancy_expert_num
-
-                w1 = torch.empty((total_expert_num, gate_out_dim + up_out_dim, gate_in_dim), dtype=dtype, device="cpu")
-
-                for i_experts in range(self.ep_n_routed_experts + self.redundancy_expert_num):
-                    w1[i_experts, 0:gate_out_dim:, :] = self.experts_gate_projs[i_experts]
-                    w1[i_experts, gate_out_dim:, :] = self.experts_up_projs[i_experts]
-
-                inter_shape, hidden_size = self.w2_list[0].shape[0], self.w2_list[0].shape[1]
-                w2 = torch._utils._flatten_dense_tensors(self.w2_list).view(len(self.w2_list), inter_shape, hidden_size)
-                if not self.quantized_weight and self.quant_method is not None:
-                    qw1, qw1_scale, qw1_zero_point = self.quant_method.quantize(w1)
-                    qw2, qw2_scale, qw2_zero_point = self.quant_method.quantize(w2)
-                    self.w1[0] = qw1
-                    self.w1[1] = qw1_scale
-                    self.w2[0] = qw2
-                    self.w2[1] = qw2_scale
-                else:
-                    self.w1[0] = self._cuda(w1)
-                    self.w2[0] = self._cuda(w2)
-                delattr(self, "w2_list")
-                delattr(self, "experts_up_projs")
-                delattr(self, "experts_gate_projs")
-
-    def _fuse_weight_scale(self):
-        with self.lock:
-            if (
-                hasattr(self, "experts_up_proj_scales")
-                and None not in self.experts_up_proj_scales
-                and None not in self.experts_gate_proj_scales
-                and None not in self.w2_scale_list
-            ):
-                gate_out_dim, gate_in_dim = self.experts_gate_proj_scales[0].shape
-                up_out_dim, up_in_dim = self.experts_up_proj_scales[0].shape
-                assert gate_in_dim == up_in_dim
-                dtype = self.experts_gate_proj_scales[0].dtype
-                total_expert_num = self.ep_n_routed_experts + self.redundancy_expert_num
-
-                w1_scale = torch.empty(
-                    (total_expert_num, gate_out_dim + up_out_dim, gate_in_dim), dtype=dtype, device="cpu"
-                )
-
-                for i_experts in range(self.ep_n_routed_experts + self.redundancy_expert_num):
-                    w1_scale[i_experts, 0:gate_out_dim:, :] = self.experts_gate_proj_scales[i_experts]
-                    w1_scale[i_experts, gate_out_dim:, :] = self.experts_up_proj_scales[i_experts]
-
-                inter_shape, hidden_size = self.w2_scale_list[0].shape[0], self.w2_scale_list[0].shape[1]
-                w2_scale = torch._utils._flatten_dense_tensors(self.w2_scale_list).view(
-                    len(self.w2_scale_list), inter_shape, hidden_size
-                )
-                self.w1[1] = self._cuda(w1_scale)
-                self.w2[1] = self._cuda(w2_scale)
-                delattr(self, "w2_scale_list")
-                delattr(self, "experts_up_proj_scales")
-                delattr(self, "experts_gate_proj_scales")
-
-    def load_hf_weights(self, weights):
-        n_expert_ep = self.ep_n_routed_experts
-        # tp to ep here
-        if self.e_score_correction_bias_name in weights:
-            self.e_score_correction_bias = self._cuda(weights[self.e_score_correction_bias_name])
-
-        for i_experts_ep in range(n_expert_ep):
-            i_experts = i_experts_ep + n_expert_ep * self.global_rank_
-            w1_weight = f"{self.weight_prefix}.{i_experts}.{self.w1_weight_name}.weight"
-            w2_weight = f"{self.weight_prefix}.{i_experts}.{self.w2_weight_name}.weight"
-            w3_weight = f"{self.weight_prefix}.{i_experts}.{self.w3_weight_name}.weight"
-            if w1_weight in weights:
-                self.experts_gate_projs[i_experts_ep] = weights[w1_weight]
-            if w3_weight in weights:
-                self.experts_up_projs[i_experts_ep] = weights[w3_weight]
-            if w2_weight in weights:
-                self.w2_list[i_experts_ep] = weights[w2_weight]
-
-        # Load weight parameters for redundant experts
-        for i, redundant_expert_id in enumerate(self.redundancy_expert_ids):
-            i_experts = redundant_expert_id
-            w1_weight = f"{self.weight_prefix}.{i_experts}.{self.w1_weight_name}.weight"
-            w2_weight = f"{self.weight_prefix}.{i_experts}.{self.w2_weight_name}.weight"
-            w3_weight = f"{self.weight_prefix}.{i_experts}.{self.w3_weight_name}.weight"
-            if w1_weight in weights:
-                self.experts_gate_projs[n_expert_ep + i] = weights[w1_weight]
-            if w3_weight in weights:
-                self.experts_up_projs[n_expert_ep + i] = weights[w3_weight]
-            if w2_weight in weights:
-                self.w2_list[n_expert_ep + i] = weights[w2_weight]
-
-        if self.quantized_weight:
-            self._load_weight_scale(weights)
-        self._fuse()
-
-    def _load_weight_scale(self, weights: Dict[str, torch.Tensor]) -> None:
-        n_expert_ep = self.ep_n_routed_experts
-        for i_experts_ep in range(n_expert_ep):
-            i_experts = i_experts_ep + n_expert_ep * self.global_rank_
-            w1_scale = f"{self.weight_prefix}.{i_experts}.{self.w1_weight_name}.{self.weight_scale_suffix}"
-            w2_scale = f"{self.weight_prefix}.{i_experts}.{self.w2_weight_name}.{self.weight_scale_suffix}"
-            w3_scale = f"{self.weight_prefix}.{i_experts}.{self.w3_weight_name}.{self.weight_scale_suffix}"
-            if w1_scale in weights:
-                self.experts_gate_proj_scales[i_experts_ep] = weights[w1_scale]
-            if w3_scale in weights:
-                self.experts_up_proj_scales[i_experts_ep] = weights[w3_scale]
-
-            if w2_scale in weights:
-                self.w2_scale_list[i_experts_ep] = weights[w2_scale]
-
-        # Load scale parameters for redundant experts
-        for i, redundant_expert_id in enumerate(self.redundancy_expert_ids):
-            i_experts = redundant_expert_id
-            w1_scale = f"{self.weight_prefix}.{i_experts}.{self.w1_weight_name}.{self.weight_scale_suffix}"
-            w2_scale = f"{self.weight_prefix}.{i_experts}.{self.w2_weight_name}.{self.weight_scale_suffix}"
-            w3_scale = f"{self.weight_prefix}.{i_experts}.{self.w3_weight_name}.{self.weight_scale_suffix}"
-            if w1_scale in weights:
-                self.experts_gate_proj_scales[n_expert_ep + i] = weights[w1_scale]
-            if w3_scale in weights:
-                self.experts_up_proj_scales[n_expert_ep + i] = weights[w3_scale]
-            if w2_scale in weights:
-                self.w2_scale_list[n_expert_ep + i] = weights[w2_scale]
-
-    def _cuda(self, cpu_tensor):
-        device_id = get_current_device_id()
-        if self.quantized_weight:
-            return cpu_tensor.contiguous().cuda(device_id)
-        return cpu_tensor.contiguous().to(self.data_type_).cuda(device_id)
-
-    def verify_load(self):
-        return self.w1 is not None and self.w2 is not None
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_tp.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_tp.py
deleted file mode 100644
index 9295fa96ae..0000000000
--- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_tp.py
+++ /dev/null
@@ -1,669 +0,0 @@
-import os
-import torch
-import threading
-from typing import Optional, Tuple, List, Dict, Any, Union
-from .base_weight import BaseWeight
-from lightllm.utils.dist_utils import get_current_rank_in_dp, get_current_device_id
-from lightllm.common.quantization import Quantcfg
-
-
-def create_tp_moe_wegiht_obj(
-    gate_proj_name: str,
-    down_proj_name: str,
-    up_proj_name: str,
-    e_score_correction_bias_name: str,
-    weight_prefix: str,
-    n_routed_experts: int,
-    num_fused_shared_experts: int,
-    split_inter_size: int,
-    data_type: torch.dtype,
-    network_config: Dict[str, Any],
-    layer_num: int,
-    quant_cfg: Quantcfg = None,
-) -> Union["FusedMoeWeightTP", "FusedAWQMARLINMoeWeightTP"]:
-    quant_method = quant_cfg.get_quant_method(layer_num, "fused_moe")
-    if quant_method is not None and quant_method.method_name == "awq_marlin":
-        return FusedAWQMARLINMoeWeightTP(
-            gate_proj_name=gate_proj_name,
-            down_proj_name=down_proj_name,
-            up_proj_name=up_proj_name,
-            e_score_correction_bias_name=e_score_correction_bias_name,
-            weight_prefix=weight_prefix,
-            n_routed_experts=n_routed_experts,
-            num_fused_shared_experts=num_fused_shared_experts,
-            split_inter_size=split_inter_size,
-            data_type=data_type,
-            network_config=network_config,
-            layer_num=layer_num,
-            quant_cfg=quant_cfg,
-        )
-    else:
-        return FusedMoeWeightTP(
-            gate_proj_name=gate_proj_name,
-            down_proj_name=down_proj_name,
-            up_proj_name=up_proj_name,
-            e_score_correction_bias_name=e_score_correction_bias_name,
-            weight_prefix=weight_prefix,
-            n_routed_experts=n_routed_experts,
-            num_fused_shared_experts=num_fused_shared_experts,
-            split_inter_size=split_inter_size,
-            data_type=data_type,
-            network_config=network_config,
-            layer_num=layer_num,
-            quant_cfg=quant_cfg,
-        )
-
-
-class FusedMoeWeightTP(BaseWeight):
-    def __init__(
-        self,
-        gate_proj_name: str,
-        down_proj_name: str,
-        up_proj_name: str,
-        e_score_correction_bias_name: str,
-        weight_prefix: str,
-        n_routed_experts: int,
-        num_fused_shared_experts: int,
-        split_inter_size: int,
-        data_type: torch.dtype,
-        network_config: Dict[str, Any],
-        layer_num: int,
-        quant_cfg: Quantcfg = None,
-    ) -> None:
-        super().__init__()
-        self.quant_method = quant_cfg.get_quant_method(layer_num, "fused_moe")
-        self.quantized_weight = quant_cfg.quantized_weight
-        if self.quant_method is not None:
-            self.weight_scale_suffix = self.quant_method.weight_scale_suffix
-            self.quant_method.is_moe = True
-        self.w1_weight_name = gate_proj_name
-        self.w2_weight_name = down_proj_name
-        self.w3_weight_name = up_proj_name
-
-        self.e_score_correction_bias_name = e_score_correction_bias_name
-        self.weight_prefix = weight_prefix
-        assert num_fused_shared_experts in [0, 1], "num_fused_shared_experts can only support 0 or 1 now."
-        self.n_routed_experts = n_routed_experts + num_fused_shared_experts
-        self.num_fused_shared_experts = num_fused_shared_experts
-        self.routed_scaling_factor = network_config.get("routed_scaling_factor", 1.0)
-        self.split_inter_size = split_inter_size
-        self.data_type_ = data_type
-        self.tp_rank_ = get_current_rank_in_dp()
-        self.experts_up_projs = [None] * self.n_routed_experts
-        self.experts_gate_projs = [None] * self.n_routed_experts
-        self.experts_up_proj_scales = [None] * self.n_routed_experts
-        self.experts_gate_proj_scales = [None] * self.n_routed_experts
-        self.e_score_correction_bias = None
-        self.w2_list = [None] * self.n_routed_experts
-        self.w2_scale_list = [None] * self.n_routed_experts
-        self.scoring_func = network_config.get("scoring_func", "softmax")
-        self.w1 = [None, None]  # weight, weight_scale
-        self.w2 = [None, None]  # weight, weight_scale
-        self.lock = threading.Lock()
-
-    def experts(self, input_tensor, router_logits, top_k, renormalize, use_grouped_topk, topk_group, num_expert_group):
-        from lightllm.common.fused_moe.topk_select import select_experts
-
-        topk_weights, topk_ids = select_experts(
-            hidden_states=input_tensor,
-            router_logits=router_logits,
-            correction_bias=self.e_score_correction_bias,
-            use_grouped_topk=use_grouped_topk,
-            top_k=top_k,
-            renormalize=renormalize,
-            topk_group=topk_group,
-            num_expert_group=num_expert_group,
-            scoring_func=self.scoring_func,
-        )
-        topk_weights.mul_(self.routed_scaling_factor)
-        if self.num_fused_shared_experts > 0:
-            pad_topk_ids = (
-                torch.arange(
-                    start=self.n_routed_experts - self.num_fused_shared_experts,
-                    end=self.n_routed_experts,
-                    step=1,
-                    dtype=topk_ids.dtype,
-                    device="cuda",
-                )
-                .view(1, self.num_fused_shared_experts)
-                .repeat(topk_ids.shape[0], 1)
-            )
-            pad_topk_weights = torch.full(
-                (topk_weights.shape[0], self.num_fused_shared_experts),
-                fill_value=1.0,
-                device="cuda",
-                dtype=topk_weights.dtype,
-            )
-
-            topk_ids = torch.cat([topk_ids, pad_topk_ids], dim=1)
-            topk_weights = torch.cat([topk_weights, pad_topk_weights], dim=1)
-
-        w1, w1_scale = self.w1
-        w2, w2_scale = self.w2
-        use_fp8_w8a8 = self.quant_method is not None
-
-        from lightllm.common.fused_moe.grouped_fused_moe import fused_experts
-
-        fused_experts(
-            hidden_states=input_tensor,
-            w1=w1,
-            w2=w2,
-            topk_weights=topk_weights,
-            topk_ids=topk_ids,
-            inplace=True,
-            use_fp8_w8a8=use_fp8_w8a8,
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-        )
-        return
-
-    def _fuse(self):
-        if self.quantized_weight:
-            self._fuse_weight_scale()
-        with self.lock:
-            if (
-                hasattr(self, "experts_up_projs")
-                and None not in self.experts_up_projs
-                and None not in self.experts_gate_projs
-                and None not in self.w2_list
-            ):
-                gate_out_dim, gate_in_dim = self.experts_gate_projs[0].shape
-                up_out_dim, up_in_dim = self.experts_up_projs[0].shape
-                assert gate_in_dim == up_in_dim
-                dtype = self.experts_gate_projs[0].dtype
-                total_expert_num = self.n_routed_experts
-
-                w1 = torch.empty((total_expert_num, gate_out_dim + up_out_dim, gate_in_dim), dtype=dtype, device="cpu")
-
-                for i_experts in range(self.n_routed_experts):
-                    w1[i_experts, 0:gate_out_dim:, :] = self.experts_gate_projs[i_experts]
-                    w1[i_experts, gate_out_dim:, :] = self.experts_up_projs[i_experts]
-
-                inter_shape, hidden_size = self.w2_list[0].shape[0], self.w2_list[0].shape[1]
-                w2 = torch._utils._flatten_dense_tensors(self.w2_list).view(len(self.w2_list), inter_shape, hidden_size)
-                if not self.quantized_weight and self.quant_method is not None:
-                    qw1, qw1_scale, qw1_zero_point = self.quant_method.quantize(w1)
-                    qw2, qw2_scale, qw2_zero_point = self.quant_method.quantize(w2)
-                    self.w1[0] = qw1
-                    self.w1[1] = qw1_scale
-                    self.w2[0] = qw2
-                    self.w2[1] = qw2_scale
-                else:
-                    self.w1[0] = self._cuda(w1)
-                    self.w2[0] = self._cuda(w2)
-                delattr(self, "w2_list")
-                delattr(self, "experts_up_projs")
-                delattr(self, "experts_gate_projs")
-
-    def _fuse_weight_scale(self):
-        with self.lock:
-            if (
-                hasattr(self, "experts_up_proj_scales")
-                and None not in self.experts_up_proj_scales
-                and None not in self.experts_gate_proj_scales
-                and None not in self.w2_scale_list
-            ):
-                gate_out_dim, gate_in_dim = self.experts_gate_proj_scales[0].shape
-                up_out_dim, up_in_dim = self.experts_up_proj_scales[0].shape
-                assert gate_in_dim == up_in_dim
-                dtype = self.experts_gate_proj_scales[0].dtype
-                total_expert_num = self.n_routed_experts
-
-                w1_scale = torch.empty(
-                    (total_expert_num, gate_out_dim + up_out_dim, gate_in_dim), dtype=dtype, device="cpu"
-                )
-
-                for i_experts in range(self.n_routed_experts):
-                    w1_scale[i_experts, 0:gate_out_dim:, :] = self.experts_gate_proj_scales[i_experts]
-                    w1_scale[i_experts, gate_out_dim:, :] = self.experts_up_proj_scales[i_experts]
-                inter_shape, hidden_size = self.w2_scale_list[0].shape[0], self.w2_scale_list[0].shape[1]
-                w2_scale = torch._utils._flatten_dense_tensors(self.w2_scale_list).view(
-                    len(self.w2_scale_list), inter_shape, hidden_size
-                )
-                self.w1[1] = self._cuda(w1_scale)
-                self.w2[1] = self._cuda(w2_scale)
-                delattr(self, "w2_scale_list")
-                delattr(self, "experts_up_proj_scales")
-                delattr(self, "experts_gate_proj_scales")
-
-    def load_hf_weights(self, weights):
-        if self.e_score_correction_bias_name in weights:
-            self.e_score_correction_bias = self._cuda(weights[self.e_score_correction_bias_name])
-        for i_experts in range(self.n_routed_experts):
-            w1_weight = f"{self.weight_prefix}.{i_experts}.{self.w1_weight_name}.weight"
-            w2_weight = f"{self.weight_prefix}.{i_experts}.{self.w2_weight_name}.weight"
-            w3_weight = f"{self.weight_prefix}.{i_experts}.{self.w3_weight_name}.weight"
-
-            if w1_weight in weights:
-                self.experts_gate_projs[i_experts] = weights[w1_weight][
-                    self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1), :
-                ]
-            if w3_weight in weights:
-                self.experts_up_projs[i_experts] = weights[w3_weight][
-                    self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1), :
-                ]
-
-            if w2_weight in weights:
-                self.w2_list[i_experts] = weights[w2_weight][
-                    :, self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1)
-                ]
-        if self.quant_method is not None:
-            self._load_weight_scale(weights)
-        self._fuse()
-
-    def _load_weight_scale(self, weights: Dict[str, torch.Tensor]) -> None:
-        block_size = 1
-        if hasattr(self.quant_method, "block_size"):
-            block_size = self.quant_method.block_size
-        for i_experts in range(self.n_routed_experts):
-            w1_scale = f"{self.weight_prefix}.{i_experts}.{self.w1_weight_name}.{self.weight_scale_suffix}"
-            w2_scale = f"{self.weight_prefix}.{i_experts}.{self.w2_weight_name}.{self.weight_scale_suffix}"
-            w3_scale = f"{self.weight_prefix}.{i_experts}.{self.w3_weight_name}.{self.weight_scale_suffix}"
-            if w1_scale in weights:
-                self.experts_gate_proj_scales[i_experts] = weights[w1_scale][
-                    self.split_inter_size
-                    // block_size
-                    * self.tp_rank_ : self.split_inter_size
-                    // block_size
-                    * (self.tp_rank_ + 1),
-                    :,
-                ]
-            if w3_scale in weights:
-                self.experts_up_proj_scales[i_experts] = weights[w3_scale][
-                    self.split_inter_size
-                    // block_size
-                    * self.tp_rank_ : self.split_inter_size
-                    // block_size
-                    * (self.tp_rank_ + 1),
-                    :,
-                ]
-
-            if w2_scale in weights:
-                self.w2_scale_list[i_experts] = weights[w2_scale][
-                    :,
-                    self.split_inter_size
-                    // block_size
-                    * self.tp_rank_ : self.split_inter_size
-                    // block_size
-                    * (self.tp_rank_ + 1),
-                ]
-
-    def _cuda(self, cpu_tensor):
-        device_id = get_current_device_id()
-        if self.quantized_weight:
-            return cpu_tensor.contiguous().cuda(device_id)
-        return cpu_tensor.contiguous().to(self.data_type_).cuda(device_id)
-
-    def verify_load(self):
-        return self.w1 is not None and self.w2 is not None
-
-
-class FusedAWQMARLINMoeWeightTP(BaseWeight):
-    def __init__(
-        self,
-        gate_proj_name: str,
-        down_proj_name: str,
-        up_proj_name: str,
-        e_score_correction_bias_name: str,
-        weight_prefix: str,
-        n_routed_experts: int,
-        num_fused_shared_experts: int,
-        split_inter_size: int,
-        data_type: torch.dtype,
-        network_config: Dict[str, Any],
-        layer_num: int,
-        quant_cfg: Quantcfg = None,
-    ) -> None:
-        super().__init__()
-        self.quant_method = quant_cfg.get_quant_method(layer_num, "fused_moe")
-        self.quantized_weight = quant_cfg.quantized_weight
-        if self.quant_method is not None:
-            self.weight_scale_suffix = self.quant_method.weight_scale_suffix
-            self.weight_zero_point_suffix = self.quant_method.weight_zero_point_suffix
-            self.quant_method.is_moe = True
-        hf_quantization_config = network_config.get("quantization_config", None)
-        self.num_bits = hf_quantization_config.get("bits", 4)
-        self.group_size = hf_quantization_config.get("group_size", 128)
-        self.pack_factor = 32 // self.num_bits
-        self.has_processed_weight = False
-        assert self.quant_method.method_name == "awq_marlin"
-
-        self.w1_weight_name = gate_proj_name
-        self.w2_weight_name = down_proj_name
-        self.w3_weight_name = up_proj_name
-
-        self.e_score_correction_bias_name = e_score_correction_bias_name
-        self.weight_prefix = weight_prefix
-        assert num_fused_shared_experts in [0, 1], "num_fused_shared_experts can only support 0 or 1 now."
-        self.n_routed_experts = n_routed_experts + num_fused_shared_experts
-        self.num_fused_shared_experts = num_fused_shared_experts
-        self.routed_scaling_factor = network_config.get("routed_scaling_factor", 1.0)
-        self.split_inter_size = split_inter_size
-        self.data_type_ = data_type
-        self.tp_rank_ = get_current_rank_in_dp()
-        self.experts_up_projs = [None] * self.n_routed_experts
-        self.experts_gate_projs = [None] * self.n_routed_experts
-        self.experts_up_proj_scales = [None] * self.n_routed_experts
-        self.experts_up_proj_zero_points = [None] * self.n_routed_experts
-        self.experts_gate_proj_scales = [None] * self.n_routed_experts
-        self.experts_gate_proj_zero_points = [None] * self.n_routed_experts
-        self.e_score_correction_bias = None
-        self.w2_list = [None] * self.n_routed_experts
-        self.w2_scale_list = [None] * self.n_routed_experts
-        self.w2_zero_point_list = [None] * self.n_routed_experts
-        self.scoring_func = network_config.get("scoring_func", "softmax")
-        self.w1 = [None, None, None]  # weight, weight_scale, zero_point
-        self.w2 = [None, None, None]  # weight, weight_scale, zero_point
-        self.lock = threading.Lock()
-
-    def experts(self, input_tensor, router_logits, top_k, renormalize, use_grouped_topk, topk_group, num_expert_group):
-        from lightllm.common.fused_moe.topk_select import select_experts
-
-        topk_weights, topk_ids = select_experts(
-            hidden_states=input_tensor,
-            router_logits=router_logits,
-            correction_bias=self.e_score_correction_bias,
-            use_grouped_topk=use_grouped_topk,
-            top_k=top_k,
-            renormalize=renormalize,
-            topk_group=topk_group,
-            num_expert_group=num_expert_group,
-            scoring_func=self.scoring_func,
-        )
-        topk_weights.mul_(self.routed_scaling_factor)
-        if self.num_fused_shared_experts > 0:
-            pad_topk_ids = (
-                torch.arange(
-                    start=self.n_routed_experts - self.num_fused_shared_experts,
-                    end=self.n_routed_experts,
-                    step=1,
-                    dtype=topk_ids.dtype,
-                    device="cuda",
-                )
-                .view(1, self.num_fused_shared_experts)
-                .repeat(topk_ids.shape[0], 1)
-            )
-            pad_topk_weights = torch.full(
-                (topk_weights.shape[0], self.num_fused_shared_experts),
-                fill_value=1.0,
-                device="cuda",
-                dtype=topk_weights.dtype,
-            )
-
-            topk_ids = torch.cat([topk_ids, pad_topk_ids], dim=1)
-            topk_weights = torch.cat([topk_weights, pad_topk_weights], dim=1)
-
-        w1, w1_scale, w1_zero_point = self.w1
-        w2, w2_scale, w2_zero_point = self.w2
-
-        from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
-
-        fused_marlin_moe(
-            input_tensor,
-            w1,
-            w2,
-            None,
-            None,
-            w1_scale,
-            w2_scale,
-            router_logits,
-            topk_weights,
-            topk_ids,
-            quant_type_id=self.quant_method.vllm_quant_type.id,
-            apply_router_weight_on_input=False,
-            global_num_experts=-1,
-            expert_map=None,
-            w1_zeros=w1_zero_point,
-            w2_zeros=w2_zero_point,
-            workspace=self.workspace,
-            inplace=True,
-        )
-
-        return
-
-    def _fuse(self):
-        self._fuse_weight()
-        self._fuse_weight_scale()
-        self._fuse_weight_zero_point()
-
-    def _fuse_weight(self):
-        with self.lock:
-            if (
-                hasattr(self, "experts_up_projs")
-                and None not in self.experts_up_projs
-                and None not in self.experts_gate_projs
-                and None not in self.w2_list
-            ):
-                gate_in_dim, gate_out_dim = self.experts_gate_projs[0].shape
-                up_in_dim, up_out_dim = self.experts_up_projs[0].shape
-                assert gate_in_dim == up_in_dim
-                total_expert_num = self.n_routed_experts
-
-                w1 = torch.empty(
-                    (total_expert_num, gate_in_dim, gate_out_dim + up_out_dim), dtype=torch.int32, device="cpu"
-                )
-
-                for i_experts in range(self.n_routed_experts):
-                    w1[i_experts, :, 0:gate_out_dim] = self.experts_gate_projs[i_experts]
-                    w1[i_experts, :, gate_out_dim:] = self.experts_up_projs[i_experts]
-
-                inter_shape, hidden_size = self.w2_list[0].shape[0], self.w2_list[0].shape[1]
-                w2 = torch._utils._flatten_dense_tensors(self.w2_list).view(len(self.w2_list), inter_shape, hidden_size)
-                self.w1[0] = self._cuda(w1)
-                self.w2[0] = self._cuda(w2)
-                delattr(self, "w2_list")
-                delattr(self, "experts_up_projs")
-                delattr(self, "experts_gate_projs")
-
-    def _fuse_weight_scale(self):
-        with self.lock:
-            if (
-                hasattr(self, "experts_up_proj_scales")
-                and None not in self.experts_up_proj_scales
-                and None not in self.experts_gate_proj_scales
-                and None not in self.w2_scale_list
-            ):
-                gate_in_dim, gate_out_dim = self.experts_gate_proj_scales[0].shape
-                up_in_dim, up_out_dim = self.experts_up_proj_scales[0].shape
-                dtype = self.experts_gate_proj_scales[0].dtype
-                assert gate_in_dim == up_in_dim
-                total_expert_num = self.n_routed_experts
-                w1_scale = torch.empty(
-                    (total_expert_num, gate_in_dim, gate_out_dim + up_out_dim), dtype=dtype, device="cpu"
-                )
-                for i_experts in range(self.n_routed_experts):
-                    w1_scale[i_experts, :, 0:gate_out_dim] = self.experts_gate_proj_scales[i_experts]
-                    w1_scale[i_experts, :, gate_out_dim:] = self.experts_up_proj_scales[i_experts]
-                inter_shape, hidden_size = self.w2_scale_list[0].shape[0], self.w2_scale_list[0].shape[1]
-                w2_scale = torch._utils._flatten_dense_tensors(self.w2_scale_list).view(
-                    len(self.w2_scale_list), inter_shape, hidden_size
-                )
-                self.w1[1] = self._cuda(w1_scale).to(self.data_type_)
-                self.w2[1] = self._cuda(w2_scale).to(self.data_type_)
-                delattr(self, "w2_scale_list")
-                delattr(self, "experts_up_proj_scales")
-                delattr(self, "experts_gate_proj_scales")
-
-    def _fuse_weight_zero_point(self):
-        with self.lock:
-            if (
-                hasattr(self, "experts_up_proj_zero_points")
-                and None not in self.experts_up_proj_zero_points
-                and None not in self.experts_gate_proj_zero_points
-                and None not in self.w2_zero_point_list
-            ):
-                gate_in_dim, gate_out_dim = self.experts_gate_proj_zero_points[0].shape
-                up_in_dim, up_out_dim = self.experts_up_proj_zero_points[0].shape
-                assert gate_in_dim == up_in_dim
-                total_expert_num = self.n_routed_experts
-                w1_zero_point = torch.empty(
-                    (total_expert_num, gate_in_dim, gate_out_dim + up_out_dim), dtype=torch.int32, device="cpu"
-                )
-                for i_experts in range(self.n_routed_experts):
-                    w1_zero_point[i_experts, :, 0:gate_out_dim] = self.experts_gate_proj_zero_points[i_experts]
-                    w1_zero_point[i_experts, :, gate_out_dim:] = self.experts_up_proj_zero_points[i_experts]
-                inter_shape, hidden_size = self.w2_zero_point_list[0].shape[0], self.w2_zero_point_list[0].shape[1]
-                w2_zero_point = torch._utils._flatten_dense_tensors(self.w2_zero_point_list).view(
-                    len(self.w2_zero_point_list), inter_shape, hidden_size
-                )
-                self.w1[2] = self._cuda(w1_zero_point)
-                self.w2[2] = self._cuda(w2_zero_point)
-                delattr(self, "w2_zero_point_list")
-                delattr(self, "experts_up_proj_zero_points")
-                delattr(self, "experts_gate_proj_zero_points")
-
-    def load_hf_weights(self, weights):
-        self._load_weight(weights)
-        self._load_weight_scale(weights)
-        self._load_weight_zero_point(weights)
-        self._fuse()
-        self._process_weight_after_loading()
-
-    def _load_weight(self, weights: Dict[str, torch.Tensor]) -> None:
-        # awq quantization weight shape: in x out
-        if self.e_score_correction_bias_name in weights:
-            self.e_score_correction_bias = self._cuda(weights[self.e_score_correction_bias_name])
-        for i_experts in range(self.n_routed_experts):
-            w1_weight = f"{self.weight_prefix}.{i_experts}.{self.w1_weight_name}.qweight"
-            w2_weight = f"{self.weight_prefix}.{i_experts}.{self.w2_weight_name}.qweight"
-            w3_weight = f"{self.weight_prefix}.{i_experts}.{self.w3_weight_name}.qweight"
-
-            if w1_weight in weights:
-                self.experts_gate_projs[i_experts] = weights[w1_weight][
-                    :, self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1)
-                ]
-            if w3_weight in weights:
-                self.experts_up_projs[i_experts] = weights[w3_weight][
-                    :, self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1)
-                ]
-
-            if w2_weight in weights:
-                self.w2_list[i_experts] = weights[w2_weight][
-                    self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1), :
-                ]
-
-    def _load_weight_scale(self, weights: Dict[str, torch.Tensor]) -> None:
-        for i_experts in range(self.n_routed_experts):
-            w1_scale = f"{self.weight_prefix}.{i_experts}.{self.w1_weight_name}.{self.weight_scale_suffix}"
-            w2_scale = f"{self.weight_prefix}.{i_experts}.{self.w2_weight_name}.{self.weight_scale_suffix}"
-            w3_scale = f"{self.weight_prefix}.{i_experts}.{self.w3_weight_name}.{self.weight_scale_suffix}"
-            split_inter_size = self.split_inter_size * self.pack_factor
-            if w1_scale in weights:
-                self.experts_gate_proj_scales[i_experts] = weights[w1_scale][
-                    :,
-                    split_inter_size * self.tp_rank_ : split_inter_size * (self.tp_rank_ + 1),
-                ]
-            if w3_scale in weights:
-                self.experts_up_proj_scales[i_experts] = weights[w3_scale][
-                    :,
-                    split_inter_size * self.tp_rank_ : split_inter_size * (self.tp_rank_ + 1),
-                ]
-
-            if w2_scale in weights:
-                self.w2_scale_list[i_experts] = weights[w2_scale][
-                    split_inter_size * self.tp_rank_ : split_inter_size * (self.tp_rank_ + 1),
-                    :,
-                ]
-
-    def _load_weight_zero_point(self, weights: Dict[str, torch.Tensor]) -> None:
-        for i_experts in range(self.n_routed_experts):
-            w1_zero_point = f"{self.weight_prefix}.{i_experts}.{self.w1_weight_name}.{self.weight_zero_point_suffix}"
-            w2_zero_point = f"{self.weight_prefix}.{i_experts}.{self.w2_weight_name}.{self.weight_zero_point_suffix}"
-            w3_zero_point = f"{self.weight_prefix}.{i_experts}.{self.w3_weight_name}.{self.weight_zero_point_suffix}"
-            if w1_zero_point in weights:
-                self.experts_gate_proj_zero_points[i_experts] = weights[w1_zero_point][
-                    :,
-                    self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1),
-                ]
-            if w3_zero_point in weights:
-                self.experts_up_proj_zero_points[i_experts] = weights[w3_zero_point][
-                    :,
-                    self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1),
-                ]
-            if w2_zero_point in weights:
-                self.w2_zero_point_list[i_experts] = weights[w2_zero_point][
-                    self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1),
-                    :,
-                ]
-
-    def _process_weight_after_loading(self):
-        with self.lock:
-            if None in self.w1 or None in self.w2 or self.has_processed_weight:
-                return
-        self.has_processed_weight = True
-        from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops
-
-        assert HAS_VLLM, "moe awq marlin quantization requires kernels of vllm"
-
-        from vllm.model_executor.layers.quantization.utils.marlin_utils import (
-            marlin_moe_permute_scales,
-            moe_awq_to_marlin_zero_points,
-            marlin_make_workspace_new,
-        )
-
-        num_experts = self.n_routed_experts
-        device = self.w1[0].device
-
-        self.w13_g_idx_sort_indices = torch.nn.Parameter(
-            torch.empty((num_experts, 0), dtype=torch.int32, device=device),
-            requires_grad=False,
-        )
-        self.w2_g_idx_sort_indices = torch.nn.Parameter(
-            torch.empty((num_experts, 0), dtype=torch.int32, device=device),
-            requires_grad=False,
-        )
-        self.w1[0] = vllm_ops.awq_marlin_moe_repack(
-            self.w1[0],
-            self.w13_g_idx_sort_indices,
-            size_k=self.w1[0].shape[1],
-            size_n=self.w1[0].shape[2] * self.pack_factor,
-            num_bits=self.num_bits,
-        )
-
-        self.w2[0] = vllm_ops.awq_marlin_moe_repack(
-            self.w2[0],
-            self.w2_g_idx_sort_indices,
-            size_k=self.w2[0].shape[1],
-            size_n=self.w2[0].shape[2] * self.pack_factor,
-            num_bits=self.num_bits,
-        )
-
-        # Why does this take the intermediate size for size_k?
-        self.w1[1] = marlin_moe_permute_scales(
-            s=self.w1[1],
-            size_k=self.split_inter_size * self.pack_factor,
-            size_n=self.w1[1].shape[2],
-            group_size=self.group_size,
-        )
-
-        self.w2[1] = marlin_moe_permute_scales(
-            s=self.w2[1],
-            size_k=self.split_inter_size * self.pack_factor,
-            size_n=self.w2[1].shape[2],
-            group_size=self.group_size,
-        )
-
-        self.w1[2] = moe_awq_to_marlin_zero_points(
-            self.w1[2],
-            size_k=self.w1[2].shape[1],
-            size_n=self.w1[2].shape[2] * self.pack_factor,
-            num_bits=self.num_bits,
-        )
-
-        self.w2[2] = moe_awq_to_marlin_zero_points(
-            self.w2[2],
-            size_k=self.w2[2].shape[1],
-            size_n=self.w2[2].shape[2] * self.pack_factor,
-            num_bits=self.num_bits,
-        )
-
-        self.workspace = marlin_make_workspace_new(device, 4)
-
-    def _cuda(self, cpu_tensor):
-        device_id = get_current_device_id()
-        if self.quantized_weight:
-            return cpu_tensor.cuda(device_id)
-        return cpu_tensor.cuda(device_id)
-
-    def verify_load(self):
-        return self.w1 is not None and self.w2 is not None
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/__init__.py b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/__init__.py
index 63605b1774..e9ae4f30ab 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/__init__.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/__init__.py
@@ -1,10 +1,5 @@
 from .mm_weight import (
-    MMWeightPack,
     MMWeightTpl,
 )
-from .mm_factory import (
-    MMWeight,
-    ROWMMWeight,
-    ROWBMMWeight,
-    COLMMWeight,
-)
+from .rowmm_weight import ROWMMWeight, KVROWNMMWeight, ROWBMMWeight
+from .colmm_weight import COLMMWeight
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/colmm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/colmm_weight.py
index 281f30f022..1a02e00d0f 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/colmm_weight.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/colmm_weight.py
@@ -1,19 +1,20 @@
 import torch
 from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import (
     MMWeightTpl,
-    DeepGemmFP8W8A8B128MMWeight,
-    AWQMMWeightTpl,
 )
 from lightllm.common.quantization import Quantcfg
 from lightllm.utils.dist_utils import get_current_device_id
 from lightllm.common.quantization.quantize_method import QuantizationMethod
 from typing import Dict, List, Optional, Union
-from .mm_slicer import ColSliceMixin, QuantizedColSliceMixin, AwqQuantizedColSliceMixin
+from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size
+from .mm_slicer import get_col_slice_mixin
 
 
-class StandardCOLMMWeight(MMWeightTpl):
+class COLMMWeight(MMWeightTpl):
     def __init__(
         self,
+        in_dim: int,
+        out_dims: Optional[Union[int, List[int]]],
         weight_names: Union[str, List[str]],
         data_type: torch.dtype,
         bias_names: Optional[Union[str, List[str]]] = None,
@@ -21,82 +22,19 @@ def __init__(
         tp_rank: int = None,
         tp_world_size: int = None,
     ) -> None:
+        self.tp_rank_ = tp_rank if tp_rank is not None else get_current_rank_in_dp()
+        self.tp_world_size_ = tp_world_size if tp_world_size is not None else get_dp_world_size()
+        in_dim = self._get_tp_dim(in_dim)
         super().__init__(
+            in_dim=in_dim,
+            out_dims=out_dims,
             weight_names=weight_names,
             data_type=data_type,
             bias_names=bias_names,
             quant_method=quant_method,
-            tp_rank=tp_rank,
-            tp_world_size=tp_world_size,
+            tp_rank=self.tp_rank_,
+            tp_world_size=self.tp_world_size_,
         )
-        self.param_slicer = ColSliceMixin(tp_rank=tp_rank, tp_world_size=tp_world_size)
-
-
-class DeepGemmFP8W8A8B128COLMMWeight(DeepGemmFP8W8A8B128MMWeight):
-    def __init__(
-        self,
-        weight_names: Union[str, List[str]],
-        data_type: torch.dtype,
-        bias_names: Optional[Union[str, List[str]]] = None,
-        quant_method: QuantizationMethod = None,
-        tp_rank: int = None,
-        tp_world_size: int = None,
-    ) -> None:
-        super().__init__(
-            weight_names=weight_names,
-            data_type=data_type,
-            bias_names=bias_names,
-            quant_method=quant_method,
-            tp_rank=tp_rank,
-            tp_world_size=tp_world_size,
-        )
-        self.param_slicer = QuantizedColSliceMixin(tp_rank=tp_rank, tp_world_size=tp_world_size)
-
-
-class AWQCOLMMWeight(AWQMMWeightTpl):
-    def __init__(
-        self,
-        weight_names: Union[str, List[str]],
-        data_type: torch.dtype,
-        bias_names: Optional[Union[str, List[str]]] = None,
-        quant_method: QuantizationMethod = None,
-        tp_rank: int = None,
-        tp_world_size: int = None,
-    ) -> None:
-        super().__init__(
-            weight_names=weight_names,
-            data_type=data_type,
-            bias_names=bias_names,
-            quant_method=quant_method,
-            tp_rank=tp_rank,
-            tp_world_size=tp_world_size,
-        )
-        # 注意这里不是错误，因为awq的weight是按inxout存的
-        self.param_slicer = AwqQuantizedColSliceMixin(tp_rank=tp_rank, tp_world_size=tp_world_size)
-
-
-class AWQMARLINCOLMMWeight(AWQCOLMMWeight):
-    def __init__(
-        self,
-        weight_names: Union[str, List[str]],
-        data_type: torch.dtype,
-        bias_names: Optional[Union[str, List[str]]] = None,
-        quant_method: QuantizationMethod = None,
-        tp_rank: int = None,
-        tp_world_size: int = None,
-    ) -> None:
-        super().__init__(
-            weight_names=weight_names,
-            data_type=data_type,
-            bias_names=bias_names,
-            quant_method=quant_method,
-            tp_rank=tp_rank,
-            tp_world_size=tp_world_size,
+        self.param_slicer = get_col_slice_mixin(
+            self.quant_method.method_name, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_
         )
-
-
-COLMM_WEIGHT_CLS_MAP = {
-    "deepgemm-fp8w8a8-b128": DeepGemmFP8W8A8B128COLMMWeight,
-    "awq": AWQCOLMMWeight,
-    "awq_marlin": AWQMARLINCOLMMWeight,
-}
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_factory.py b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_factory.py
deleted file mode 100644
index 464de84413..0000000000
--- a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_factory.py
+++ /dev/null
@@ -1,90 +0,0 @@
-from lightllm.common.quantization import Quantcfg
-from lightllm.common.quantization.quantize_method import QuantizationMethod
-from typing import Type, Union, Dict
-from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import (
-    MMWeightTpl,
-    BMMWeightTpl,
-)
-from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.rowmm_weight import (
-    StandardROWMMWeight,
-    UnquantizedROWBMMWeight,
-    ROWMM_WEIGHT_CLS_MAP,
-)
-from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.colmm_weight import (
-    StandardCOLMMWeight,
-    COLMM_WEIGHT_CLS_MAP,
-)
-
-
-class MMWeight:
-    def __new__(cls, **kwargs):
-        """
-        weight_names,
-        data_type,
-        bias_names,
-        quant_cfg,
-        layer_num,
-        name,
-        tp_rank,
-        tp_world_size,
-        ...
-        该类主要是通过重载 __new__ 为对应的mm权重绑定量化方法，其他参数都是透传。
-        """
-
-        quant_cfg = kwargs.pop("quant_cfg", None)
-        layer_num_ = kwargs.pop("layer_num", None)
-        name = kwargs.pop("name", None)
-        quant_method, quantized_weight = cls._get_quant_method(quant_cfg, layer_num_, name)
-        # quantized_weight 本身是用来标识权重本身在文件中是否是以量化后的形式存储，
-        # 现在不再使用该参数，是否量化由后续的加载过程自动识别。
-        kwargs["quant_method"] = quant_method
-        mmcls = cls._get_mmcls(quant_method)
-        return mmcls(**kwargs)
-
-    @classmethod
-    def _get_quant_method(cls, quant_cfg: Quantcfg, layer_num_: int, name: str) -> QuantizationMethod:
-        if quant_cfg is None:
-            return None, False
-        quant_method: QuantizationMethod = quant_cfg.get_quant_method(layer_num_, name)
-        if quant_method is None:
-            return None, False
-        quant_method.hf_quantization_config = quant_cfg.hf_quantization_config
-        quantized_weight = quant_cfg.quantized_weight
-        return quant_method, quantized_weight
-
-    @classmethod
-    def _get_mmcls(cls, quant_method: QuantizationMethod) -> Type[Union[MMWeightTpl, BMMWeightTpl]]:
-        raise NotImplementedError("Subclasses must implement _get_mmcls method")
-
-
-class ROWMMWeight(MMWeight):
-    @classmethod
-    def _get_mmcls(cls, quant_method: QuantizationMethod):
-        if quant_method is None:
-            return StandardROWMMWeight
-
-        return ROWMM_WEIGHT_CLS_MAP.get(
-            quant_method.method_name,
-            StandardROWMMWeight,
-        )
-
-
-class ROWBMMWeight(MMWeight):
-    @classmethod
-    def _get_mmcls(cls, quant_method: QuantizationMethod):
-        if quant_method is None:
-            return UnquantizedROWBMMWeight
-        else:
-            # TODO: Implement more quantization weight
-            raise NotImplementedError("ROWBMMWeight is not implemented")
-
-
-class COLMMWeight(MMWeight):
-    @classmethod
-    def _get_mmcls(cls, quant_method: QuantizationMethod):
-        if quant_method is None:
-            return StandardCOLMMWeight
-        return COLMM_WEIGHT_CLS_MAP.get(
-            quant_method.method_name,
-            StandardCOLMMWeight,
-        )
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_slicer.py b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_slicer.py
index e3ef5b0ea3..ddbf98a866 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_slicer.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_slicer.py
@@ -1,5 +1,5 @@
 import torch
-from typing import Optional
+from typing import Optional, Tuple
 from abc import ABC, abstractmethod
 from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size
 
@@ -7,9 +7,12 @@
 class SliceMixinBase(ABC):
     """切片操作的Mixin基类"""
 
-    def __init__(self, tp_rank: int = None, tp_world_size: int = None):
+    def __init__(self, tp_rank: int = None, tp_world_size: int = None, repeat_times: int = 1):
         self.tp_rank_ = tp_rank if tp_rank is not None else get_current_rank_in_dp()
         self.tp_world_size_ = tp_world_size if tp_world_size is not None else get_dp_world_size()
+        # this param is used to slice the weight when tp_world_size_ is divisible by the kv_head_num
+        # for example, if tp_world_size_ is 8 and kv_head_num is 4, then repeat_times_ is 2
+        self.repeat_times_ = repeat_times
 
     @abstractmethod
     def _slice_weight(self, weight: torch.Tensor):
@@ -19,10 +22,16 @@ def _slice_weight(self, weight: torch.Tensor):
     def _slice_bias(self, bias):
         pass
 
+    def _get_slice_start_end(self, size: int) -> Tuple[int, int]:
+        tp_size = size * self.repeat_times_ // self.tp_world_size_
+        start = tp_size * (self.tp_rank_ // self.repeat_times_)
+        end = start + tp_size
+        return start, end
+
 
 class SliceMixinTpl(SliceMixinBase):
-    def __init__(self, tp_rank: int = None, tp_world_size: int = None):
-        super().__init__(tp_rank, tp_world_size)
+    def __init__(self, tp_rank: int = None, tp_world_size: int = None, repeat_times: int = 1):
+        super().__init__(tp_rank, tp_world_size, repeat_times)
 
     def _slice_weight(self, weight: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError("slice_weight must implement this method")
@@ -40,95 +49,117 @@ def _slice_weight_zero_point(self, weight_zero_point: torch.Tensor) -> torch.Ten
 # 默认weight 的shape是 outxin，这也是目前最通用的约定。
 # 所以row-wise是沿着dim=0进行切分，col-wise是沿着dim=1进行切分。
 class RowSliceMixin(SliceMixinTpl):
-    def __init__(self, tp_rank: int = None, tp_world_size: int = None):
-        super().__init__(tp_rank, tp_world_size)
+    def __init__(self, tp_rank: int = None, tp_world_size: int = None, repeat_times: int = 1):
+        super().__init__(tp_rank, tp_world_size, repeat_times)
 
     def _slice_weight(self, weight: torch.Tensor) -> torch.Tensor:
-        assert weight.shape[0] % self.tp_world_size_ == 0, f"tp slice error {weight.shape[0]} % {self.tp_world_size_}"
-        tp_size = weight.shape[0] // self.tp_world_size_
-        return weight[tp_size * self.tp_rank_ : tp_size * (self.tp_rank_ + 1)]
+        assert (
+            weight.shape[0] * self.repeat_times_ % self.tp_world_size_ == 0
+        ), f"tp slice error {weight.shape[0] * self.repeat_times_} % {self.tp_world_size_}"
+        start, end = self._get_slice_start_end(weight.shape[0])
+        return weight[start:end, :]
 
     def _slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
-        assert bias.shape[0] % self.tp_world_size_ == 0, f"tp slice error {bias.shape[0]} % {self.tp_world_size_}"
-        tp_size = bias.shape[0] // self.tp_world_size_
-        return bias[tp_size * self.tp_rank_ : tp_size * (self.tp_rank_ + 1)]
+        assert (
+            bias.shape[0] * self.repeat_times_ % self.tp_world_size_ == 0
+        ), f"tp slice error {bias.shape[0] * self.repeat_times_} % {self.tp_world_size_}"
+        start, end = self._get_slice_start_end(bias.shape[0])
+        return bias[start:end]
 
 
 # 量化切片默认实现方式是group-wise的量化，所以weight_scale 和weight_zero_point ndims跟weight一样。
 # 后续按需要，扩展per-tensor、per-channel的量化方式。
 class QuantizedRowSliceMixin(RowSliceMixin):
-    def __init__(self, tp_rank: int = None, tp_world_size: int = None):
-        super().__init__(tp_rank, tp_world_size)
+    def __init__(self, tp_rank: int = None, tp_world_size: int = None, repeat_times: int = 1):
+        super().__init__(tp_rank, tp_world_size, repeat_times)
 
     def _slice_weight_scale(self, weight_scale: torch.Tensor) -> torch.Tensor:
         assert (
             weight_scale.shape[0] % self.tp_world_size_ == 0
         ), f"tp slice error {weight_scale.shape[0]} % {self.tp_world_size_}"
-        tp_size = weight_scale.shape[0] // self.tp_world_size_
-        scale_start = tp_size * self.tp_rank_
-        scale_end = tp_size * (self.tp_rank_ + 1)
-        return weight_scale[scale_start:scale_end]
+        start, end = self._get_slice_start_end(weight_scale.shape[0])
+        return weight_scale[start:end]
 
     def _slice_weight_zero_point(self, weight_zero_point: torch.Tensor) -> torch.Tensor:
         assert (
             weight_zero_point.shape[0] % self.tp_world_size_ == 0
         ), f"tp slice error {weight_zero_point.shape[0]} % {self.tp_world_size_}"
-        tp_size = weight_zero_point.shape[0] // self.tp_world_size_
-        zero_point_start = tp_size * self.tp_rank_
-        zero_point_end = tp_size * (self.tp_rank_ + 1)
-        return weight_zero_point[zero_point_start:zero_point_end]
+        start, end = self._get_slice_start_end(weight_zero_point.shape[0])
+        return weight_zero_point[start:end]
 
 
 class ColSliceMixin(SliceMixinTpl):
-    def __init__(self, tp_rank: int = None, tp_world_size: int = None):
-        super().__init__(tp_rank, tp_world_size)
+    def __init__(self, tp_rank: int = None, tp_world_size: int = None, repeat_times: int = 1):
+        super().__init__(tp_rank, tp_world_size, repeat_times)
 
     def _slice_weight(self, weight: torch.Tensor) -> torch.Tensor:
-        assert weight.shape[1] % self.tp_world_size_ == 0, f"tp slice error {weight.shape[1]} % {self.tp_world_size_}"
-        tp_size = weight.shape[1] // self.tp_world_size_
-        return weight[:, tp_size * self.tp_rank_ : tp_size * (self.tp_rank_ + 1)]
+        assert (
+            weight.shape[1] * self.repeat_times_ % self.tp_world_size_ == 0
+        ), f"tp slice error {weight.shape[1] * self.repeat_times_ } % {self.tp_world_size_}"
+        start, end = self._get_slice_start_end(weight.shape[1])
+        return weight[:, start:end]
 
     def _slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
-        return bias / self.tp_world_size_
+        return bias / self.tp_world_size_ * self.repeat_times_
 
 
 class QuantizedColSliceMixin(ColSliceMixin):
-    def __init__(self, tp_rank: int = None, tp_world_size: int = None):
-        super().__init__(tp_rank, tp_world_size)
+    def __init__(self, tp_rank: int = None, tp_world_size: int = None, repeat_times: int = 1):
+        super().__init__(tp_rank, tp_world_size, repeat_times)
 
     def _slice_weight_scale(self, weight_scale: torch.Tensor) -> torch.Tensor:
         assert (
-            weight_scale.shape[1] % self.tp_world_size_ == 0
-        ), f"tp slice error {weight_scale.shape[1]} % {self.tp_world_size_}"
-        tp_size = weight_scale.shape[1] // self.tp_world_size_
-        scale_start = tp_size * self.tp_rank_
-        scale_end = tp_size * (self.tp_rank_ + 1)
-        return weight_scale[:, scale_start:scale_end]
+            weight_scale.shape[1] * self.repeat_times_ % self.tp_world_size_ == 0
+        ), f"tp slice error {weight_scale.shape[1] * self.repeat_times_ } % {self.tp_world_size_}"
+        start, end = self._get_slice_start_end(weight_scale.shape[1])
+        return weight_scale[:, start:end]
 
     def _slice_weight_zero_point(self, weight_zero_point: torch.Tensor) -> torch.Tensor:
         assert (
-            weight_zero_point.shape[1] % self.tp_world_size_ == 0
-        ), f"tp slice error {weight_zero_point.shape[1]} % {self.tp_world_size_}"
-        tp_size = weight_zero_point.shape[1] // self.tp_world_size_
-        zero_point_start = tp_size * self.tp_rank_
-        zero_point_end = tp_size * (self.tp_rank_ + 1)
-        return weight_zero_point[:, zero_point_start:zero_point_end]
+            weight_zero_point.shape[1] * self.repeat_times_ % self.tp_world_size_ == 0
+        ), f"tp slice error {weight_zero_point.shape[1] * self.repeat_times_ } % {self.tp_world_size_}"
+        start, end = self._get_slice_start_end(weight_zero_point.shape[1])
+        return weight_zero_point[:, start:end]
 
 
 # awq 的量化权重是inxout存储格式，需要定制实现。
 class AwqQuantizedRowSliceMixin(QuantizedColSliceMixin):
-    def __init__(self, tp_rank: int = None, tp_world_size: int = None):
-        super().__init__(tp_rank, tp_world_size)
+    def __init__(self, tp_rank: int = None, tp_world_size: int = None, repeat_times: int = 1):
+        super().__init__(tp_rank, tp_world_size, repeat_times)
 
     def _slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
-        assert bias.shape[0] % self.tp_world_size_ == 0, f"tp slice error {bias.shape[0]} % {self.tp_world_size_}"
-        tp_size = bias.shape[0] // self.tp_world_size_
-        return bias[tp_size * self.tp_rank_ : tp_size * (self.tp_rank_ + 1)]
+        assert (
+            bias.shape[0] * self.repeat_times_ % self.tp_world_size_ == 0
+        ), f"tp slice error {bias.shape[0] * self.repeat_times_ } % {self.tp_world_size_}"
+        start, end = self._get_slice_start_end(bias.shape[0])
+        return bias[start:end]
 
 
 class AwqQuantizedColSliceMixin(QuantizedRowSliceMixin):
-    def __init__(self, tp_rank: int = None, tp_world_size: int = None):
-        super().__init__(tp_rank, tp_world_size)
+    def __init__(self, tp_rank: int = None, tp_world_size: int = None, repeat_times: int = 1):
+        super().__init__(tp_rank, tp_world_size, repeat_times)
 
     def _slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
-        return bias / self.tp_world_size_
+        return bias / self.tp_world_size_ * self.repeat_times_
+
+
+def get_row_slice_mixin(
+    quant_method_name: str, tp_rank: int = None, tp_world_size: int = None, repeat_times: int = 1
+) -> SliceMixinTpl:
+    if quant_method_name.startswith("awq"):
+        return AwqQuantizedRowSliceMixin(tp_rank, tp_world_size, repeat_times)
+    elif quant_method_name == "none":
+        return RowSliceMixin(tp_rank, tp_world_size, repeat_times)
+    else:
+        return QuantizedRowSliceMixin(tp_rank, tp_world_size, repeat_times)
+
+
+def get_col_slice_mixin(
+    quant_method_name: str, tp_rank: int = None, tp_world_size: int = None, repeat_times: int = 1
+) -> SliceMixinTpl:
+    if quant_method_name.startswith("awq"):
+        return AwqQuantizedColSliceMixin(tp_rank, tp_world_size, repeat_times)
+    elif quant_method_name == "none":
+        return ColSliceMixin(tp_rank, tp_world_size, repeat_times)
+    else:
+        return QuantizedColSliceMixin(tp_rank, tp_world_size, repeat_times)
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
index 7391454da0..56aa322b49 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
@@ -5,9 +5,10 @@
 from dataclasses import dataclass
 from typing import Optional, Tuple, List, Dict, Union, Type
 from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
-from lightllm.common.quantization.quantize_method import QuantizationMethod
+from lightllm.common.quantization.quantize_method import QuantizationMethod, WeightPack
 from lightllm.common.basemodel.layer_weights.meta_weights.base_weight import BaseWeightTpl
 from lightllm.common.quantization import Quantcfg
+from lightllm.common.quantization.no_quant import NoQuantization
 from lightllm.utils.dist_utils import get_current_device_id
 from lightllm.utils.log_utils import init_logger
 from .mm_slicer import SliceMixinTpl
@@ -15,53 +16,11 @@
 logger = init_logger(__name__)
 
 
-@dataclass
-class MMWeightPack:
-    weight: Optional[torch.Tensor] = None
-    bias: Optional[torch.Tensor] = None
-    weight_scale: Optional[torch.Tensor] = None
-    weight_zero_point: Optional[torch.Tensor] = None
-
-    has_bias: bool = False
-    has_weight_scale: bool = False
-    has_weight_zero_point: bool = False
-
-    def is_ready(self) -> bool:
-        return (
-            self.weight is not None
-            and (not self.has_bias or (self.has_bias and self.bias is not None))
-            and (not self.has_weight_scale or (self.has_weight_scale and self.weight_scale is not None))
-            and (not self.has_weight_zero_point or (self.has_weight_zero_point and self.weight_zero_point is not None))
-        )
-
-    def ready_for_fused_merge(self) -> bool:
-        """
-        判断权重是否满足可以和其他权重进行融合cat的条件，因为可能权重是量化和非量化后的权重，所以复杂一些。
-        """
-        weight_ready = self.weight is not None and self.weight.dtype in [
-            torch.bfloat16,
-            torch.float16,
-            torch.float32,
-            torch.float64,
-        ]
-        bias_ready = (self.has_bias and self.bias is not None) or (not self.has_bias)
-        if weight_ready and bias_ready:
-            return True
-        else:
-            return self.is_ready()
-
-    def is_load_finished(self):
-        return (
-            (self.is_ready() and self.weight.is_cuda)
-            and ((self.has_bias and self.bias.is_cuda) or (not self.has_bias))
-            and ((self.has_weight_scale and self.weight_scale.is_cuda) or (not self.has_weight_scale))
-            and ((self.has_weight_zero_point and self.weight_zero_point.is_cuda) or (not self.has_weight_zero_point))
-        )
-
-
 class MMWeightTpl(BaseWeightTpl):
     def __init__(
         self,
+        in_dim: int,
+        out_dims: Optional[Union[int, List[int]]],
         weight_names: Union[str, List[str]],
         bias_names: Optional[Union[str, List[str]]],
         data_type: torch.dtype,
@@ -70,7 +29,14 @@ def __init__(
         tp_world_size: int = None,
     ) -> None:
         super().__init__(tp_rank, tp_world_size, data_type)
-        self.lock = threading.Lock()
+
+        self.in_dim = in_dim
+        if isinstance(out_dims, int):
+            out_dims = [out_dims]
+        self.out_dims = out_dims
+        self.cusum_out_dims = [0]
+        for out_dim in out_dims[:-1]:
+            self.cusum_out_dims.append(self.cusum_out_dims[-1] + out_dim)
 
         if isinstance(weight_names, str):
             weight_names = [weight_names]
@@ -82,64 +48,23 @@ def __init__(
             if bias_names[0] is None:
                 bias_names = None
 
-        if quant_method is not None:
-            has_weight_scale = quant_method.has_weight_scale
-            has_weight_zero_point = quant_method.has_weight_zero_point
-        else:
-            has_weight_scale = False
-            has_weight_zero_point = False
-
         # 同时存在 weight_names 和 quanted_weight_names 是为了兼容在线和离线两种加载方案
         self.weight_names = weight_names
-
         self.bias_names = bias_names
-        has_bias = self.bias_names is not None
-
-        self.gen_weight_quant_param_names(quant_method=quant_method)
-        self.quant_method = quant_method
-        self.sub_child_mm_params: List[MMWeightPack] = [
-            MMWeightPack(
-                has_bias=has_bias,
-                has_weight_scale=has_weight_scale,
-                has_weight_zero_point=has_weight_zero_point,
-            )
-            for _ in range(len(weight_names))
-        ]
-        self.mm_param: MMWeightPack = MMWeightPack(
-            has_bias=has_bias,
-            has_weight_scale=has_weight_scale,
-            has_weight_zero_point=has_weight_zero_point,
-        )
+        self.quant_method: QuantizationMethod = NoQuantization() if quant_method is None else quant_method
         self.param_slicer: SliceMixinTpl = None
-
-        self.weight_fused_dim = 0
-        self.bias_fused_dim = 0
-        self.weight_scale_and_zero_point_fused_dim = 0
-
-        self.load_finished: bool = False
+        self._create_weight()
+        self.gen_weight_quant_param_names(quant_method=quant_method)
 
     def mm(
         self, input_tensor: torch.Tensor, out: Optional[torch.Tensor] = None, use_custom_tensor_mananger: bool = True
     ) -> torch.Tensor:
-        if self.quant_method is not None:
-            return self.quant_method.apply(
-                input_tensor, self.mm_param, out, use_custom_tensor_mananger=use_custom_tensor_mananger
-            )
-        if out is None:
-            shape = (input_tensor.shape[0], self.mm_param.weight.shape[1])
-            dtype = input_tensor.dtype
-            device = input_tensor.device
-            if use_custom_tensor_mananger:
-                out = g_cache_manager.alloc_tensor(shape, dtype, device=device)
-            else:
-                out = torch.empty(shape, dtype=dtype, device=device)
-        if self.mm_param.bias is None:
-            return torch.mm(input_tensor, self.mm_param.weight, out=out)
-        return torch.addmm(self.mm_param.bias, input_tensor, self.mm_param.weight, out=out)
+        return self.quant_method.apply(
+            input_tensor, self.mm_param, out, use_custom_tensor_mananger=use_custom_tensor_mananger, bias=self.bias
+        )
 
     def gen_weight_quant_param_names(self, quant_method: Optional[QuantizationMethod]):
         if quant_method is None:
-            self.quanted_weight_names = None
             self.weight_zero_point_names = None
             self.weight_scale_names = None
             return
@@ -160,9 +85,7 @@ def gen_weight_quant_param_names(self, quant_method: Optional[QuantizationMethod
                 quanted_weight_names.append(weight_name)
 
         if len(quanted_weight_names) != 0:
-            self.quanted_weight_names = quanted_weight_names
-        else:
-            self.quanted_weight_names = None
+            self.weight_names = quanted_weight_names
 
         if len(weight_scale_names) != 0:
             self.weight_scale_names = weight_scale_names
@@ -176,16 +99,10 @@ def gen_weight_quant_param_names(self, quant_method: Optional[QuantizationMethod
         return
 
     def load_hf_weights(self, weights):
-        if self.mm_param.is_load_finished():
-            return
 
         for sub_child_index, param_name in enumerate(self.weight_names):
             self._load_weight(param_name=param_name, weights=weights, sub_child_index=sub_child_index)
 
-        if self.quanted_weight_names is not None:
-            for sub_child_index, param_name in enumerate(self.quanted_weight_names):
-                self._load_weight(param_name=param_name, weights=weights, sub_child_index=sub_child_index)
-
         if self.bias_names is not None:
             for sub_child_index, param_name in enumerate(self.bias_names):
                 self._load_bias(param_name=param_name, weights=weights, sub_child_index=sub_child_index)
@@ -196,51 +113,16 @@ def load_hf_weights(self, weights):
             for sub_child_index, param_name in enumerate(self.weight_zero_point_names):
                 self._load_weight_zero_point(param_name=param_name, weights=weights, sub_child_index=sub_child_index)
 
-        with self.lock:
-            # 如果需要fused的请求，全部ok了以后进行merge操作。, all([]) 竟然返回是True, 需要len(self.sub_child_mm_params) > 0 的额外判断。
-            if len(self.sub_child_mm_params) > 0 and all(e.ready_for_fused_merge() for e in self.sub_child_mm_params):
-                self._fuse_weights()
-                self.sub_child_mm_params.clear()
-
-            # 在线量化操作
-            if (
-                self.quant_method is not None
-                and self.mm_param.weight is not None
-                and self.quant_method.weight_need_quanted(self.mm_param.weight)
-                and self.load_finished is False
-            ):
-                logger.info(f"online quant weight names: {self.weight_names}")
-                quantized_weight, weight_scale, weight_zero_point = self.quant_method.quantize(
-                    self.mm_param.weight.cuda(get_current_device_id())
-                )
-                self.mm_param.weight = quantized_weight
-                self.mm_param.weight_scale = weight_scale
-                self.mm_param.weight_zero_point = weight_zero_point
-
-            # repack 操作
-            if (
-                self.quant_method is not None
-                and self.mm_param.is_ready()
-                and self.quant_method.params_need_repack()
-                and self.load_finished is False
-            ):
-                (
-                    self.mm_param.weight,
-                    self.mm_param.weight_scale,
-                    self.mm_param.weight_zero_point,
-                ) = self.quant_method.params_repack(
-                    weight=self.mm_param.weight,
-                    weight_scale=self.mm_param.weight_scale,
-                    weight_zero_point=self.mm_param.weight_zero_point,
-                    dtype_type=self.data_type_,
-                )
-
-            if self.mm_param.is_ready() and self.load_finished is False:
-                self._to_gpu_device()
-                self.load_finished = True
-
-    def verify_load(self) -> bool:
-        return self.mm_param.is_ready()
+    def _create_weight(self):
+        self.bias = None
+        if self.bias_names is not None:
+            self.bias = torch.empty(sum(self.out_dims), dtype=self.data_type_).cuda(get_current_device_id())
+            self.bias._load_ok = [False] * len(self.bias_names)
+        self.mm_param: WeightPack = self.quant_method.create_weight(
+            in_dim=self.in_dim, out_dim=sum(self.out_dims), dtype=self.data_type_, device_id=get_current_device_id()
+        )
+        self.mm_param.initialize_load_status(len(self.weight_names))
+        return
 
     # 执行顺序
     def _load_weight(
@@ -248,7 +130,16 @@ def _load_weight(
     ) -> None:
         if param_name in weights:
             weight = self.param_slicer._slice_weight(weights[param_name])
-            self.sub_child_mm_params[sub_child_index].weight = weight
+            start_idx = self.cusum_out_dims[sub_child_index]
+            if self.quant_method.weight_need_quanted(weight):
+                self.quant_method.quantize(weight, self.mm_param, offset=start_idx)
+                # weight_scale and zero_point will be computed during online quantization.
+                # so we set them to True here.
+                self.mm_param.load_ok[sub_child_index][1] = True
+                self.mm_param.load_ok[sub_child_index][2] = True
+            else:
+                self.quant_method.load_weight(weight, self.mm_param, start_idx)
+            self.mm_param.load_ok[sub_child_index][0] = True
         return
 
     def _load_bias(
@@ -256,7 +147,10 @@ def _load_bias(
     ) -> None:
         if param_name in weights:
             bias = self.param_slicer._slice_bias(weights[param_name])
-            self.sub_child_mm_params[sub_child_index].bias = bias
+            start_idx = self.cusum_out_dims[sub_child_index]
+            end_idx = start_idx + bias.shape[0]
+            self.bias[start_idx:end_idx].copy_(bias)
+            self.bias._load_ok[sub_child_index] = True
         return
 
     def _load_weight_scale(
@@ -264,7 +158,9 @@ def _load_weight_scale(
     ) -> None:
         if param_name in weights:
             weight_scale = self.param_slicer._slice_weight_scale(weights[param_name])
-            self.sub_child_mm_params[sub_child_index].weight_scale = weight_scale
+            start_idx = self.cusum_out_dims[sub_child_index]
+            self.quant_method.load_weight_scale(weight_scale, self.mm_param, start_idx)
+            self.mm_param.load_ok[sub_child_index][1] = True
         return
 
     def _load_weight_zero_point(
@@ -272,102 +168,75 @@ def _load_weight_zero_point(
     ) -> None:
         if param_name in weights:
             weight_zero_point = self.param_slicer._slice_weight_zero_point(weights[param_name])
-            self.sub_child_mm_params[sub_child_index].weight_zero_point = weight_zero_point
+            start_idx = self.cusum_out_dims[sub_child_index]
+            self.quant_method.load_weight_zero_point(weight_zero_point, self.mm_param, start_idx)
+            self.mm_param.load_ok[sub_child_index][2] = True
         return
 
-    # weight merge
-    def _fuse_weights(self) -> None:
-        need_merge = len(self.sub_child_mm_params) > 1
-        if self.mm_param.weight is None and all(p.weight is not None for p in self.sub_child_mm_params):
-            if need_merge:
-                weight = torch.cat([p.weight for p in self.sub_child_mm_params], dim=self.weight_fused_dim)
-            else:
-                weight = self.sub_child_mm_params[0].weight
-
-            # 快速删除，防止占用显存过久
-            for p in self.sub_child_mm_params:
-                p.weight = None
-
-            self.mm_param.weight = weight
-
-        if (
-            self.mm_param.has_bias
-            and self.mm_param.bias is None
-            and all(p.bias is not None for p in self.sub_child_mm_params)
-        ):
-            if need_merge:
-                bias = torch.cat([p.bias for p in self.sub_child_mm_params], dim=self.bias_fused_dim)
-            else:
-                bias = self.sub_child_mm_params[0].bias
+    def verify_load(self):
+        mm_param_load_ok = all(all(load_ok_list) for load_ok_list in self.mm_param.load_ok)
+        bias_load_ok = True if self.bias is None else all(self.bias._load_ok)
+        if not (mm_param_load_ok and bias_load_ok):
+            logger.warning(f"mm_param_load_ok: {self.mm_param.load_ok}, bias_load_ok: {self.bias}")
+        return mm_param_load_ok and bias_load_ok
 
-            # 快速删除，防止占用显存过久
-            for p in self.sub_child_mm_params:
-                p.bias = None
-
-            self.mm_param.bias = bias
-
-        if self.mm_param.weight_scale is None and all(p.weight_scale is not None for p in self.sub_child_mm_params):
-            if need_merge:
-                weight_scale = torch.cat(
-                    [p.weight_scale for p in self.sub_child_mm_params], dim=self.weight_scale_and_zero_point_fused_dim
-                )
-            else:
-                weight_scale = self.sub_child_mm_params[0].weight_scale
+    def _get_tp_dim(self, dim: int) -> int:
+        assert (
+            dim % self.tp_world_size_ == 0
+        ), f"dim must be divisible by tp_world_size_, but found: {dim} % {self.tp_world_size_}"
+        return dim // self.tp_world_size_
 
-            # 快速删除，防止占用显存过久
-            for p in self.sub_child_mm_params:
-                p.weight_scale = None
 
-            self.mm_param.weight_scale = weight_scale
-
-        if self.mm_param.weight_zero_point is None and all(
-            p.weight_zero_point is not None for p in self.sub_child_mm_params
-        ):
-            if need_merge:
-                weight_zero_point = torch.cat(
-                    [p.weight_zero_point for p in self.sub_child_mm_params],
-                    dim=self.weight_scale_and_zero_point_fused_dim,
-                )
-            else:
-                weight_zero_point = self.sub_child_mm_params[0].weight_zero_point
-
-            # 快速删除，防止占用显存过久
-            for p in self.sub_child_mm_params:
-                p.weight_zero_point = None
-
-            self.mm_param.weight_zero_point = weight_zero_point
+class BMMWeightTpl(BaseWeightTpl):
+    def __init__(
+        self,
+        dim0: int,
+        dim1: int,
+        dim2: int,
+        weight_names: Union[str, List[str]],
+        data_type: torch.dtype,
+        bias_names: Optional[Union[str, List[str]]] = None,
+        quant_method: QuantizationMethod = None,
+        tp_rank: int = None,
+        tp_world_size: int = None,
+    ) -> None:
+        super().__init__(tp_rank, tp_world_size, data_type)
+        if isinstance(weight_names, str):
+            weight_names = [weight_names]
+        self.weight_names = weight_names
+        self.bias_names = bias_names
+        assert bias_names is None, "bmm not support bias"
+        if isinstance(bias_names, list):
+            assert all(bias_name is None for bias_name in bias_names), "bmm not support bias"
+        assert quant_method is None, "bmm not support quantized weight"
+        self.quant_method = quant_method
+        self.dim0 = dim0
+        self.dim1 = dim1
+        self.dim2 = dim2
+        self._create_weight()
         return
 
-    def _to_gpu_device(self) -> None:
-        if self.mm_param.weight is not None:
-            if self.quant_method is not None:
-                self.mm_param.weight = self.mm_param.weight.cuda(get_current_device_id())
-            else:
-                # 让 k dim 更连续，大多数split k 算法的算子可能能更快
-                self.mm_param.weight = (
-                    self.mm_param.weight.to(self.data_type_).cuda(get_current_device_id()).transpose(0, 1)
-                )
-        if self.mm_param.weight_scale is not None:
-            self.mm_param.weight_scale = self.mm_param.weight_scale.cuda(get_current_device_id())
-        if self.mm_param.weight_zero_point is not None:
-            self.mm_param.weight_zero_point = self.mm_param.weight_zero_point.cuda(get_current_device_id())
-        if self.mm_param.bias is not None:
-            # TODO 是不是所有的bias都需要转换为全局设置的数据类型吗，会不会影响精度
-            self.mm_param.bias = self.mm_param.bias.to(self.data_type_).cuda(get_current_device_id())
+    def _create_weight(self):
+        self.weight = torch.empty(self.dim0, self.dim1, self.dim2, dtype=self.data_type_).cuda(get_current_device_id())
+        self.weight._load_ok = False
         return
 
+    def load_hf_weights(self, weights: Dict[str, torch.Tensor]):
+        for weight_name in self.weight_names:
+            if weight_name in weights:
+                weight = self.param_slicer._slice_weight(weights[weight_name])
+                self.weight.copy_(weight)
+                self.weight._load_ok = True
+        return
 
-class BMMWeightTpl(MMWeightTpl):
-    def mm(
-        self, input_tensor: torch.Tensor, out: Optional[torch.Tensor] = None, use_custom_tensor_mananger: bool = True
-    ) -> torch.Tensor:
-        raise RuntimeError("use bmm not mm")
+    def verify_load(self):
+        return self.weight._load_ok
 
     def bmm(
         self, input_tensor: torch.Tensor, out: Optional[torch.Tensor] = None, use_custom_tensor_mananger: bool = True
     ) -> torch.Tensor:
         # 目前 bmm 不支持量化运算操作
-        fpweight = self.mm_param.weight
+        fpweight = self.weight
         if out is None:
             shape = (input_tensor.shape[0], input_tensor.shape[1], fpweight.shape[2])
             dtype = input_tensor.dtype
@@ -376,90 +245,4 @@ def bmm(
                 out = g_cache_manager.alloc_tensor(shape, dtype, device=device)
             else:
                 out = torch.empty(shape, dtype=dtype, device=device)
-        if self.mm_param.bias is None:
-            return torch.bmm(input_tensor, fpweight, out=out)
-        return torch.addbmm(self.mm_param.bias, input_tensor, fpweight, out=out)
-
-    def _to_gpu_device(self) -> None:
-        if self.mm_param.weight is not None:
-            if self.quant_method is not None:
-                self.mm_param.weight = self.mm_param.weight.cuda(get_current_device_id())
-            else:
-                # bmm 不需要 transpose 操作
-                self.mm_param.weight = self.mm_param.weight.to(self.data_type_).cuda(get_current_device_id())
-        if self.mm_param.weight_scale is not None:
-            self.mm_param.weight_scale = self.mm_param.weight_scale.cuda(get_current_device_id())
-        if self.mm_param.weight_zero_point is not None:
-            self.mm_param.weight_zero_point = self.mm_param.weight_zero_point.cuda(get_current_device_id())
-        if self.mm_param.bias is not None:
-            # TODO 是不是所有的bias都需要转换为全局设置的数据类型吗，会不会影响精度
-            self.mm_param.bias = self.mm_param.bias.to(self.data_type_).cuda(get_current_device_id())
-        return
-
-
-class DeepGemmFP8W8A8B128MMWeight(MMWeightTpl):
-    def __init__(
-        self,
-        weight_names: Union[str, List[str]],
-        data_type: torch.dtype,
-        bias_names: Optional[Union[str, List[str]]] = None,
-        quant_method: QuantizationMethod = None,
-        tp_rank: int = None,
-        tp_world_size: int = None,
-    ) -> None:
-        super().__init__(
-            weight_names=weight_names,
-            bias_names=bias_names,
-            data_type=data_type,
-            quant_method=quant_method,
-            tp_rank=tp_rank,
-            tp_world_size=tp_world_size,
-        )
-
-    def _to_gpu_device(self) -> None:
-        if self.mm_param.weight is not None:
-            self.mm_param.weight = self.mm_param.weight.cuda(get_current_device_id()).transpose(0, 1)
-        if self.mm_param.weight_scale is not None:
-            self.mm_param.weight_scale = self.mm_param.weight_scale.cuda(get_current_device_id()).transpose(0, 1)
-
-        assert self.mm_param.has_weight_zero_point is False
-
-        if self.mm_param.bias is not None:
-            # TODO 是不是所有的bias都需要转换为全局设置的数据类型吗，会不会影响精度
-            self.mm_param.bias = self.mm_param.bias.to(self.data_type_).cuda(get_current_device_id())
-        return
-
-
-class AWQMMWeightTpl(MMWeightTpl):
-    def __init__(
-        self,
-        weight_names: Union[str, List[str]],
-        bias_names: Optional[Union[str, List[str]]] = None,
-        data_type: torch.dtype = None,
-        quant_method: QuantizationMethod = None,
-        tp_rank: int = None,
-        tp_world_size: int = None,
-    ) -> None:
-        super().__init__(
-            weight_names=weight_names,
-            bias_names=bias_names,
-            data_type=data_type,
-            quant_method=quant_method,
-            tp_rank=tp_rank,
-            tp_world_size=tp_world_size,
-        )
-        self.weight_fused_dim = 1
-        self.bias_fused_dim = 0
-        self.weight_scale_and_zero_point_fused_dim = 1
-
-    def _to_gpu_device(self) -> None:
-        if self.mm_param.weight is not None:
-            self.mm_param.weight = self.mm_param.weight.cuda(get_current_device_id())
-        if self.mm_param.weight_scale is not None:
-            self.mm_param.weight_scale = self.mm_param.weight_scale.to(self.data_type_).cuda(get_current_device_id())
-        if self.mm_param.weight_zero_point is not None:
-            self.mm_param.weight_zero_point = self.mm_param.weight_zero_point.cuda(get_current_device_id())
-        if self.mm_param.bias is not None:
-            # TODO 是不是所有的bias都需要转换为全局设置的数据类型吗，会不会影响精度
-            self.mm_param.bias = self.mm_param.bias.to(self.data_type_).cuda(get_current_device_id())
-        return
+        return torch.bmm(input_tensor, fpweight, out=out)
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/rowmm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/rowmm_weight.py
index 0eebdc74d2..30a699bb68 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/rowmm_weight.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/rowmm_weight.py
@@ -1,20 +1,18 @@
 import torch
-from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import (
-    MMWeightTpl,
-    DeepGemmFP8W8A8B128MMWeight,
-    AWQMMWeightTpl,
-    BMMWeightTpl,
-)
+from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import MMWeightTpl, BMMWeightTpl
 from lightllm.common.quantization import Quantcfg
 from lightllm.utils.dist_utils import get_current_device_id
 from lightllm.common.quantization.quantize_method import QuantizationMethod
 from typing import Dict, List, Optional, Union
-from .mm_slicer import RowSliceMixin, QuantizedRowSliceMixin, AwqQuantizedRowSliceMixin
+from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size
+from .mm_slicer import get_row_slice_mixin
 
 
-class StandardROWMMWeight(MMWeightTpl):
+class ROWMMWeight(MMWeightTpl):
     def __init__(
         self,
+        in_dim: int,
+        out_dims: Optional[Union[int, List[int]]],
         weight_names: Union[str, List[str]],
         data_type: torch.dtype,
         bias_names: Optional[Union[str, List[str]]] = None,
@@ -22,42 +20,30 @@ def __init__(
         tp_rank: int = None,
         tp_world_size: int = None,
     ) -> None:
+        self.tp_rank_ = tp_rank if tp_rank is not None else get_current_rank_in_dp()
+        self.tp_world_size_ = tp_world_size if tp_world_size is not None else get_dp_world_size()
+        out_dims = [self._get_tp_dim(out_dim) for out_dim in out_dims]
         super().__init__(
+            in_dim=in_dim,
+            out_dims=out_dims,
             weight_names=weight_names,
             bias_names=bias_names,
             data_type=data_type,
             quant_method=quant_method,
-            tp_rank=tp_rank,
-            tp_world_size=tp_world_size,
+            tp_rank=self.tp_rank_,
+            tp_world_size=self.tp_world_size_,
         )
-        self.param_slicer = RowSliceMixin(tp_rank=tp_rank, tp_world_size=tp_world_size)
-
-
-class DeepGemmFP8W8A8B128ROWMMWeight(DeepGemmFP8W8A8B128MMWeight):
-    def __init__(
-        self,
-        weight_names: Union[str, List[str]],
-        data_type: torch.dtype,
-        bias_names: Optional[Union[str, List[str]]] = None,
-        quant_method: QuantizationMethod = None,
-        tp_rank: int = None,
-        tp_world_size: int = None,
-    ) -> None:
-        super().__init__(
-            weight_names=weight_names,
-            data_type=data_type,
-            bias_names=bias_names,
-            quant_method=quant_method,
-            tp_rank=tp_rank,
-            tp_world_size=tp_world_size,
+        self.param_slicer = get_row_slice_mixin(
+            self.quant_method.method_name, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_
         )
-        self.param_slicer = QuantizedRowSliceMixin(tp_rank=tp_rank, tp_world_size=tp_world_size)
-        return
 
 
-class UnquantizedROWBMMWeight(BMMWeightTpl):
+class KVROWNMMWeight(MMWeightTpl):
     def __init__(
         self,
+        in_dim: int,
+        kv_head_num: int,
+        head_dim: int,
         weight_names: Union[str, List[str]],
         data_type: torch.dtype,
         bias_names: Optional[Union[str, List[str]]] = None,
@@ -65,42 +51,53 @@ def __init__(
         tp_rank: int = None,
         tp_world_size: int = None,
     ) -> None:
-        super().__init__(
-            weight_names=weight_names,
-            data_type=data_type,
-            bias_names=bias_names,
-            quant_method=quant_method,
-            tp_rank=tp_rank,
-            tp_world_size=tp_world_size,
+        self.tp_rank_ = tp_rank if tp_rank is not None else get_current_rank_in_dp()
+        self.tp_world_size_ = tp_world_size if tp_world_size is not None else get_dp_world_size()
+        self.repeat_times = 1
+        assert kv_head_num % self.tp_world_size_ == 0 or self.tp_world_size_ % kv_head_num == 0, (
+            f"kv_head_num must be divisible by tp_world_size_ or "
+            f"tp_world_size_ must be divisible by kv_head_num, "
+            f"but found: {kv_head_num} % {self.tp_world_size_}"
         )
-        self.param_slicer = RowSliceMixin(tp_rank=tp_rank, tp_world_size=tp_world_size)
-
-
-class AWQROWMMWeight(AWQMMWeightTpl):
-    def __init__(
-        self,
-        weight_names: Union[str, List[str]],
-        data_type: torch.dtype,
-        bias_names: Optional[Union[str, List[str]]] = None,
-        quant_method: QuantizationMethod = None,
-        tp_rank: int = None,
-        tp_world_size: int = None,
-    ) -> None:
+        kv_hidden_size = self._get_tp_padded_head_num(kv_head_num) * head_dim
+        out_dims = [kv_hidden_size, kv_hidden_size]
         super().__init__(
+            in_dim=in_dim,
+            out_dims=out_dims,
             weight_names=weight_names,
             data_type=data_type,
             bias_names=bias_names,
             quant_method=quant_method,
-            tp_rank=tp_rank,
-            tp_world_size=tp_world_size,
+            tp_rank=self.tp_rank_,
+            tp_world_size=self.tp_world_size_,
+        )
+        self.param_slicer = get_row_slice_mixin(
+            self.quant_method.method_name,
+            tp_rank=self.tp_rank_,
+            tp_world_size=self.tp_world_size_,
+            repeat_times=self.repeat_times,
         )
 
-        self.param_slicer = AwqQuantizedRowSliceMixin(tp_rank=tp_rank, tp_world_size=tp_world_size)
+    def _get_tp_padded_head_num(self, head_num: int):
+        if head_num % self.tp_world_size_ == 0:
+            return head_num // self.tp_world_size_
+        elif self.tp_world_size_ % head_num == 0:
+            self.repeat_times = self.tp_world_size_ // head_num
+            return self.repeat_times * head_num // self.tp_world_size_
+        else:
+            raise ValueError(
+                f"head_num must be divisible by tp_world_size_ or "
+                f"tp_world_size_ must be divisible by head_num, "
+                f"but found: {head_num} % {self.tp_world_size_}"
+            )
 
 
-class AWQMARLINROWMMWeight(AWQROWMMWeight):
+class ROWBMMWeight(BMMWeightTpl):
     def __init__(
         self,
+        dim0: int,
+        dim1: int,
+        dim2: int,
         weight_names: Union[str, List[str]],
         data_type: torch.dtype,
         bias_names: Optional[Union[str, List[str]]] = None,
@@ -108,18 +105,23 @@ def __init__(
         tp_rank: int = None,
         tp_world_size: int = None,
     ) -> None:
+        self.tp_rank_ = tp_rank if tp_rank is not None else get_current_rank_in_dp()
+        self.tp_world_size_ = tp_world_size if tp_world_size is not None else get_dp_world_size()
+        assert (
+            dim0 % self.tp_world_size_ == 0
+        ), f"dim0 of bmm must be divisible by tp_world_size_, but found: {dim0} % {self.tp_world_size_}"
+        dim0 = dim0 // self.tp_world_size_
         super().__init__(
+            dim0=dim0,
+            dim1=dim1,
+            dim2=dim2,
             weight_names=weight_names,
-            data_type=data_type,
             bias_names=bias_names,
+            data_type=data_type,
             quant_method=quant_method,
-            tp_rank=tp_rank,
-            tp_world_size=tp_world_size,
+            tp_rank=self.tp_rank_,
+            tp_world_size=self.tp_world_size_,
+        )
+        self.param_slicer = get_row_slice_mixin(
+            quant_method_name="none", tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_
         )
-
-
-ROWMM_WEIGHT_CLS_MAP = {
-    "deepgemm-fp8w8a8-b128": DeepGemmFP8W8A8B128ROWMMWeight,
-    "awq": AWQROWMMWeight,
-    "awq_marlin": AWQMARLINROWMMWeight,
-}
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py
index 5a595bff61..d4717386b3 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py
@@ -1,143 +1,240 @@
 import torch
-from typing import Optional
+from typing import Optional, Dict
 from .base_weight import BaseWeightTpl
-from lightllm.utils.dist_utils import get_current_device_id
-from lightllm.common.basemodel.triton_kernel.rmsnorm import rmsnorm_forward
-from lightllm.common.basemodel.triton_kernel.layernorm import layernorm_forward
-from lightllm.utils.log_utils import init_logger
+from lightllm.utils.dist_utils import get_current_device_id, get_current_rank_in_dp, get_dp_world_size
+from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward
+from lightllm.common.basemodel.triton_kernel.norm.layernorm import layernorm_forward
+from lightllm.common.basemodel.triton_kernel.norm.qk_norm import qk_rmsnorm_forward
+from .platform_op import PlatformAwareOp
 
-logger = init_logger(__name__)
 
-
-class _NormWeight(BaseWeightTpl):
-    def __init__(self, weight_name, data_type, bias_name=None):
+class RMSNormWeight(BaseWeightTpl, PlatformAwareOp):
+    def __init__(self, dim: int, weight_name: str, data_type: torch.dtype, bias_name: str = None):
         super().__init__()
+        self.dim = dim
         self.weight_name = weight_name
-        self.bias_name = bias_name
         self.data_type_ = data_type
-        self.weight: torch.Tensor = None
-        self.bias: Optional[torch.Tensor] = None
+        assert bias_name is None, "RMSNormWeight does not have bias"
+        self._create_weight()
+
+    def _create_weight(self):
+        self.weight: torch.Tensor = torch.empty(self.dim, dtype=self.data_type_, device=self.device_id_)
+        self.weight.load_ok = False
+
+    def load_hf_weights(self, weights: Dict[str, torch.Tensor]):
+        if self.weight_name in weights:
+            self.weight.copy_(weights[self.weight_name])
+            self.weight.load_ok = True
 
     def verify_load(self):
-        load_ok = True
-        # Verify weight. The weight must be not None.
-        load_ok = load_ok and self.weight is not None
-        # Verify bias. If bias_name is set, it must be not None.
-        if self.bias_name is not None:
-            load_ok = load_ok and self.bias is not None
-        return load_ok
-
-    def rmsnorm_forward(
+        return self.weight.load_ok
+
+    def _native_forward(
+        self, input: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        assert input.ndim == 2 and self.weight.ndim == 1
+        assert input.shape[-1] == self.dim, f"Expected hidden_size to be {self.dim}, but found: {input.shape[-1]}"
+        x = input.to(torch.float32)
+        x_var = x
+        variance = x_var.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + eps)
+        x = (x * self.weight).to(self.data_type_)
+        if out is not None:
+            out.copy_(x)
+            return out
+        return x
+
+    def _triton_forward(
         self, input: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
     ) -> torch.Tensor:
-        assert input.ndim in [2, 3] and self.weight.ndim == 1
-        assert self.bias is None
+        assert (
+            input.ndim in [2, 3] and self.weight.ndim == 1
+        ), f"input.ndim: {input.ndim} != 2 or weight.ndim: {self.weight.ndim} != 1"
         if out is None:
             out = alloc_func(input.shape, dtype=input.dtype, device=input.device)
         return rmsnorm_forward(x=input, weight=self.weight, eps=eps, out=out)
 
-    def layernorm_forward(
+    def _cuda_forward(
         self, input: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
     ) -> torch.Tensor:
-        assert input.ndim == 2 and self.weight.ndim == 1
-        assert self.bias is not None
-
-        _tout = layernorm_forward(x=input, weight=self.weight, bias=self.bias, eps=eps)
-        if out is None:
-            return _tout
-        else:
-            out.copy_(_tout)
-            return out
+        # only triton implementation is supported for rmsnorm on cuda platform
+        return self._triton_forward(input=input, eps=eps, out=out, alloc_func=alloc_func)
 
+    def _musa_forward(
+        self, input: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        # triton implementation is supported by musa.
+        return self._triton_forward(input=input, eps=eps, out=out, alloc_func=alloc_func)
 
-class NoTpNormWeight(_NormWeight):
-    def __init__(self, weight_name, data_type, bias_name=None):
-        super().__init__(weight_name=weight_name, data_type=data_type, bias_name=bias_name)
-        self.tp_world_size_ = 1
-        self.tp_rank_ = 0
+    def __call__(
+        self, input: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        return self._forward(input=input, eps=eps, out=out, alloc_func=alloc_func)
 
-    def load_hf_weights(self, weights):
-        if self.weight_name in weights and self.weight is None:
-            self.weight = weights[self.weight_name].to(self.data_type_).cuda(get_current_device_id())
-        if self.bias_name in weights and self.bias is None:
-            self.bias = weights[self.bias_name].to(self.data_type_).cuda(get_current_device_id())
 
+class LayerNormWeight(BaseWeightTpl, PlatformAwareOp):
+    def __init__(self, dim: int, weight_name: str, data_type: torch.dtype, bias_name: str = None):
+        super().__init__()
+        self.dim = dim
+        self.weight_name = weight_name
+        self.bias_name = bias_name
+        self.data_type_ = data_type
+        self._create_weight()
+
+    def _create_weight(self):
+        self.weight: torch.Tensor = torch.empty(self.dim, dtype=self.data_type_, device=self.device_id_)
+        self.bias: torch.Tensor = torch.empty(self.dim, dtype=self.data_type_, device=self.device_id_)
+        self.weight.load_ok = False
+        self.bias.load_ok = False
+
+    def load_hf_weights(self, weights: Dict[str, torch.Tensor]):
+        if self.weight_name in weights:
+            self.weight.copy_(weights[self.weight_name])
+            self.weight.load_ok = True
+        if self.bias_name in weights:
+            self.bias.copy_(weights[self.bias_name])
+            self.bias.load_ok = True
 
-class NoTpGEMMANormWeight(_NormWeight):
-    def __init__(self, weight_name, data_type, bias_name=None):
-        super().__init__(weight_name, data_type, bias_name)
-        assert self.bias_name is None
-        self.tp_world_size_ = 1
-        self.tp_rank_ = 0
+    def verify_load(self):
+        return self.weight.load_ok and self.bias.load_ok
 
-    def load_hf_weights(self, weights):
-        if self.weight_name in weights and self.weight is None:
-            self.weight = (weights[self.weight_name] + 1).to(self.data_type_).cuda(get_current_device_id())
+    def _native_forward(
+        self, input: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        assert input.ndim == 2 and self.weight.ndim == 1
+        assert input.shape[-1] == self.dim, f"Expected hidden_size to be {self.dim}, but found: {input.shape[-1]}"
+        x = torch.nn.functional.layer_norm(
+            input, normalized_shape=[self.dim], weight=self.weight, bias=self.bias, eps=eps
+        )
+        if out is not None:
+            out.copy_(x.to(self.data_type_))
+            return out
+        return x.to(self.data_type_)
 
+    def _triton_forward(
+        self, input: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        assert input.ndim == 2 and self.weight.ndim == 1
+        if out is None:
+            out = alloc_func(input.shape, dtype=input.dtype, device=input.device)
+        return layernorm_forward(x=input, weight=self.weight, bias=self.bias, eps=eps, out=out)
 
-class TpVitPadNormWeight(_NormWeight):
-    def __init__(self, weight_name, data_type, head_num: int, bias_name=None):
-        super().__init__(weight_name, data_type, bias_name)
-        self.head_num = head_num
+    def _cuda_forward(
+        self, input: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        # only triton implementation is supported for layernorm on cuda platform
+        return self._triton_forward(input=input, eps=eps, out=out, alloc_func=alloc_func)
 
-    def _pad_tensor_param(self, weight: torch.Tensor):
-        assert weight.ndim == 1
-        hidden_size = weight.shape[0]
-        head_dim = hidden_size // self.head_num
-        assert hidden_size % self.head_num == 0
+    def _musa_forward(
+        self, input: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        # triton implementation is supported by musa.
+        return self._triton_forward(input=input, eps=eps, out=out, alloc_func=alloc_func)
 
-        if self.head_num % self.tp_world_size_ == 0:
-            return weight
+    def __call__(
+        self, input: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None, alloc_func=torch.empty
+    ) -> torch.Tensor:
+        return self._forward(input=input, eps=eps, out=out, alloc_func=alloc_func)
+
+
+class TpRMSNormWeight(RMSNormWeight):
+    def __init__(self, dim: int, weight_name: str, data_type: torch.dtype, bias_name: str = None):
+        super().__init__(dim=dim, weight_name=weight_name, data_type=data_type, bias_name=bias_name)
+        self.tp_world_size_ = get_dp_world_size()
+        self.tp_rank_ = get_current_rank_in_dp()
+        self.dim = self._get_tp_padded_dim(dim=dim)
+        self.repeat_times_ = 1
+
+    def _get_tp_padded_dim(self, dim: int):
+        """
+        Get the padded dimension for the weight.
+        1. if dim is divisible by tp_world_size_, return dim
+        2. if dim is greater than tp_world_size_, return (dim + tp_world_size_ - 1) // tp_world_size_ * tp_world_size_
+        3. if dim is less than tp_world_size_, assert tp_world_size_ is divisible by dim, and return dim
+        """
+        if dim % self.tp_world_size_ == 0:
+            return dim // self.tp_world_size_
+
+        if dim > self.tp_world_size_:
+            return (dim + self.tp_world_size_ - 1) // self.tp_world_size_ * self.tp_world_size_
         else:
-            logger.warning(f"padding {self.weight_name} weights in TpVitPadNormWeight")
-            pad_head_num = self.tp_world_size_ - (self.head_num % self.tp_world_size_)
-            pad_dims = pad_head_num * head_dim
-            weight = torch.nn.functional.pad(weight, (0, pad_dims), mode="constant", value=0.0)
-            return weight
+            assert (
+                self.tp_world_size_ % dim == 0
+            ), f"tp_world_size_ must be divisible by dim, but found: {self.tp_world_size_} % {dim}"
+            self.repeat_times_ = self.tp_world_size_ // dim
+            return dim * self.repeat_times_ // self.tp_world_size_
 
     def load_hf_weights(self, weights):
         if self.weight_name in weights and self.weight is None:
             t_weight = weights[self.weight_name]
-            t_weight = self._pad_tensor_param(t_weight)
-            new_hidden_size = t_weight.shape[0]
-            split_n_embed = new_hidden_size // self.tp_world_size_
-            assert new_hidden_size % self.tp_world_size_ == 0
+            hidden_size = t_weight.shape[0]
+            split_hidden_size = hidden_size // self.tp_world_size_
 
-            start = split_n_embed * self.tp_rank_
-            end = split_n_embed * (self.tp_rank_ + 1)
+            start = split_hidden_size * self.tp_rank_ // self.repeat_times_
+            end = min(split_hidden_size * (self.tp_rank_ + 1) // self.repeat_times_, hidden_size)
 
-            self.weight = t_weight[start:end].to(self.data_type_).cuda(get_current_device_id())
+            self.weight[:, end - start].copy_(t_weight[start:end].to(self.data_type_))
+            # the padding part is zero
+            self.weight[:, end:].zero_()
+            self.weight.load_ok = True
 
-        if self.bias_name in weights and self.bias is None:
-            t_bias = weights[self.bias_name]
-            t_bias = self._pad_tensor_param(t_bias)
-            new_hidden_size = t_bias.shape[0]
-            split_n_embed = new_hidden_size // self.tp_world_size_
-            assert new_hidden_size % self.tp_world_size_ == 0
 
-            start = split_n_embed * self.tp_rank_
-            end = split_n_embed * (self.tp_rank_ + 1)
+class NoTpGEMMANormWeight(RMSNormWeight):
+    def __init__(self, dim: int, weight_name: str, data_type: torch.dtype, bias_name: str = None):
+        super().__init__(dim=dim, weight_name=weight_name, data_type=data_type, bias_name=bias_name)
+        self.tp_world_size_ = 1
+        self.tp_rank_ = 0
 
-            self.bias = t_bias[start:end].to(self.data_type_).cuda(get_current_device_id())
+    def load_hf_weights(self, weights: Dict[str, torch.Tensor]):
+        if self.weight_name in weights:
+            self.weight.copy_(weights[self.weight_name])
+            self.weight += 1
+            self.load_cnt += 1
 
 
-class TpHeadNormWeight(_NormWeight):
-    def __init__(self, weight_name, data_type, bias_name=None):
-        super().__init__(weight_name, data_type, bias_name)
+class QKRMSNORMWeight(RMSNormWeight):
+    def __init__(self, dim: int, weight_name: str, data_type: torch.dtype, bias_name: str = None):
+        super().__init__(dim=dim, weight_name=weight_name, data_type=data_type, bias_name=bias_name)
+        self.tp_world_size_ = 1
+        self.tp_rank_ = 0
 
-    def load_hf_weights(self, weights):
-        if self.weight_name in weights and self.weight is None:
-            t_weight = weights[self.weight_name]
-            start_head_index, end_head_index = self._get_head_tp_split_params(weight=t_weight)
-            self.weight: torch.Tensor = (
-                t_weight[start_head_index:end_head_index].to(self.data_type_).cuda(get_current_device_id())
-            )
-            assert self.weight.ndim == 2
-
-        if self.bias_name in weights and self.bias is None:
-            t_bias = weights[self.bias_name]
-            start_head_index, end_head_index = self._get_head_tp_split_params(weight=t_bias)
-            self.bias: torch.Tensor = (
-                t_bias[start_head_index:end_head_index].to(self.data_type_).cuda(get_current_device_id())
-            )
-            assert self.bias.ndim == 2
+    def _native_forward(
+        self,
+        input: torch.Tensor,
+        eps: float,
+    ) -> None:
+        assert input.ndim == 2 and self.weight.ndim == 1
+        assert input.shape[-1] == self.dim, f"Expected hidden_size to be {self.dim}, but found: {input.shape[-1]}"
+        head_dim = self.weight.shape[0]
+        x = input.to(torch.float32)
+        x = x.view(-1, head_dim)
+        x_var = x
+        variance = x_var.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + eps)
+        x = (x * self.weight).to(self.data_type_)
+        x = x.view(-1, input.shape[-1])
+        input.copy_(x)
+        return
+
+    def _triton_forward(self, input: torch.Tensor, eps: float) -> torch.Tensor:
+        assert input.ndim == 2 and self.weight.ndim == 1
+        return qk_rmsnorm_forward(x=input, weight=self.weight, eps=eps)
+
+    def _cuda_forward(
+        self,
+        input: torch.Tensor,
+        eps: float,
+    ) -> None:
+        self._triton_forward(input=input, eps=eps)
+        return
+
+    def _musa_forward(self, input: torch.Tensor, eps: float) -> torch.Tensor:
+        # musa implementation is supported by musa triton on musa platform
+        return self._triton_forward(input=input, eps=eps)
+
+    def __call__(
+        self,
+        input: torch.Tensor,
+        eps: float,
+    ) -> None:
+        return self._forward(input=input, eps=eps)
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/platform_op.py b/lightllm/common/basemodel/layer_weights/meta_weights/platform_op.py
new file mode 100644
index 0000000000..1ba1610fc9
--- /dev/null
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/platform_op.py
@@ -0,0 +1,71 @@
+import torch
+from abc import ABC, abstractmethod
+from typing import Optional, Callable, Any
+from lightllm.utils.device_utils import get_platform, Platform
+from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class PlatformAwareOp(ABC):
+    """
+    platform aware op base class,
+    automatically route to the corresponding implementation method according to the platform.
+    """
+
+    def __init__(self):
+        args = get_env_start_args()
+        self.platform = get_platform(args.hardware_platform)
+        self.enable_torch_fallback = args.enable_torch_fallback
+        self.enable_triton_fallback = args.enable_triton_fallback
+        self._forward = self._route_forward()
+
+    def _route_forward(self) -> Callable:
+
+        method_name_map = {
+            Platform.CUDA: "_cuda_forward",
+            Platform.ASCEND: "_ascend_forward",
+            Platform.CAMBRICON: "_cambricon_forward",
+            Platform.MUSA: "_musa_forward",
+            Platform.ROCM: "_rocm_forward",
+            Platform.CPU: "_cpu_forward",
+        }
+
+        method_name = method_name_map.get(self.platform)
+        if method_name and hasattr(self, method_name):
+            method = getattr(self, method_name)
+            if callable(method):
+                return method
+
+        if self.enable_triton_fallback:
+            if hasattr(self, "_triton_forward"):
+                return self._triton_forward
+            logger.warning(
+                f"No triton implementation found for {self.__class__.__name__} on {self.platform.name} platform. "
+                f"Please implement {self.__class__.__name__}_{self.platform.name}_triton_forward method, "
+                f"or set --enable_torch_fallback to use default implementation."
+            )
+
+        if self.enable_torch_fallback:
+            return self._native_forward
+
+        # if no implementation found, raise error
+        raise NotImplementedError(
+            f"No implementation found for {self.__class__.__name__} on {self.platform.name} platform. "
+            f"Please implement {self.__class__.__name__}_{self.platform.name}_forward method, "
+            f"or set --enable_torch_fallback to use default implementation."
+        )
+
+    @abstractmethod
+    def _native_forward(self, *args, **kwargs) -> Any:
+        raise NotImplementedError("default forward must implement this method")
+
+    @abstractmethod
+    def _cuda_forward(self, *args, **kwargs) -> Any:
+        raise NotImplementedError("cuda forward must implement this method")
+
+    # Since Triton may be compatible with all hardware platforms in the future,
+    # so provide triton implementation as a fallback for all hardware platforms
+    def _triton_forward(self, *args, **kwargs) -> Any:
+        raise NotImplementedError("triton forward must implement this method")
diff --git a/lightllm/common/basemodel/layer_weights/pre_and_post_layer_weight.py b/lightllm/common/basemodel/layer_weights/pre_and_post_layer_weight.py
index 19eb67017d..8c81fd5bcc 100644
--- a/lightllm/common/basemodel/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/common/basemodel/layer_weights/pre_and_post_layer_weight.py
@@ -3,11 +3,10 @@
 
 
 class PreAndPostLayerWeight(BaseLayerWeight):
-    def __init__(self, data_type, network_config, mode):
+    def __init__(self, data_type, network_config):
         super().__init__()
         self.data_type_ = data_type
         self.network_config_ = network_config
-        self.mode = mode
         self.init_static_params()
         return
 
diff --git a/lightllm/common/basemodel/layer_weights/transformer_layer_weight.py b/lightllm/common/basemodel/layer_weights/transformer_layer_weight.py
index 97bc762370..c6ce1049f0 100644
--- a/lightllm/common/basemodel/layer_weights/transformer_layer_weight.py
+++ b/lightllm/common/basemodel/layer_weights/transformer_layer_weight.py
@@ -4,17 +4,17 @@
 from .base_layer_weight import BaseLayerWeight
 from .meta_weights import BaseWeight, MMWeightTpl
 from lightllm.utils.log_utils import init_logger
+from lightllm.common.quantization import Quantcfg
 
 logger = init_logger(__name__)
 
 
 class TransformerLayerWeight(BaseLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode, quant_cfg):
+    def __init__(self, layer_num, data_type, network_config, quant_cfg):
         super().__init__()
         self.layer_num_ = layer_num
         self.data_type_ = data_type
         self.network_config_ = network_config
-        self.mode = mode
         self.quant_cfg = quant_cfg
         self._parse_config()
         self._init_weight_names()
@@ -41,3 +41,6 @@ def load_hf_weights(self, weights):
                     attr.load_hf_weights(weights)
             elif isinstance(attr, BaseWeight):
                 attr.load_hf_weights(weights)
+
+    def get_quant_method(self, name):
+        return self.quant_cfg.get_quant_method(self.layer_num_, name)
diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py
index 3d77a3ae4c..a8b2616418 100644
--- a/lightllm/common/basemodel/prefill_cuda_graph.py
+++ b/lightllm/common/basemodel/prefill_cuda_graph.py
@@ -168,7 +168,6 @@ def warmup(self, model):
             model_input = ModelInput(
                 batch_size=1,
                 total_token_num=total_token_num,
-                max_len_in_batch=total_token_num,
                 max_q_seq_len=total_token_num,
                 max_kv_seq_len=total_token_num,
                 max_cache_len=0,
@@ -229,7 +228,6 @@ def warmup_overlap(self, model):
                 micro_batch = ModelInput(
                     batch_size=1,
                     total_token_num=total_token_num,
-                    max_len_in_batch=total_token_num,
                     max_q_seq_len=total_token_num,
                     max_kv_seq_len=total_token_num,
                     max_cache_len=0,
diff --git a/lightllm/models/bloom/triton_kernel/__init__.py b/lightllm/common/basemodel/triton_kernel/alibi_att/__init__.py
similarity index 100%
rename from lightllm/models/bloom/triton_kernel/__init__.py
rename to lightllm/common/basemodel/triton_kernel/alibi_att/__init__.py
diff --git a/lightllm/models/bloom/triton_kernel/context_flashattention_nopad.py b/lightllm/common/basemodel/triton_kernel/alibi_att/context_flashattention_nopad.py
similarity index 100%
rename from lightllm/models/bloom/triton_kernel/context_flashattention_nopad.py
rename to lightllm/common/basemodel/triton_kernel/alibi_att/context_flashattention_nopad.py
diff --git a/lightllm/models/bloom/triton_kernel/token_attention_nopad_att1.py b/lightllm/common/basemodel/triton_kernel/alibi_att/token_attention_nopad_att1.py
similarity index 100%
rename from lightllm/models/bloom/triton_kernel/token_attention_nopad_att1.py
rename to lightllm/common/basemodel/triton_kernel/alibi_att/token_attention_nopad_att1.py
diff --git a/lightllm/models/bloom/triton_kernel/token_attention_nopad_reduceV.py b/lightllm/common/basemodel/triton_kernel/alibi_att/token_attention_nopad_reduceV.py
similarity index 100%
rename from lightllm/models/bloom/triton_kernel/token_attention_nopad_reduceV.py
rename to lightllm/common/basemodel/triton_kernel/alibi_att/token_attention_nopad_reduceV.py
diff --git a/lightllm/models/bloom/triton_kernel/token_attention_nopad_softmax.py b/lightllm/common/basemodel/triton_kernel/alibi_att/token_attention_nopad_softmax.py
similarity index 77%
rename from lightllm/models/bloom/triton_kernel/token_attention_nopad_softmax.py
rename to lightllm/common/basemodel/triton_kernel/alibi_att/token_attention_nopad_softmax.py
index 25af80fabf..adf97735f6 100644
--- a/lightllm/models/bloom/triton_kernel/token_attention_nopad_softmax.py
+++ b/lightllm/common/basemodel/triton_kernel/alibi_att/token_attention_nopad_softmax.py
@@ -6,11 +6,15 @@
 
 @triton.jit
 def _fwd_kernel_token_softmax(
-    Logics, B_Start_Loc, B_Seqlen,
+    Logics,
+    B_Start_Loc,
+    B_Seqlen,
     Prob_Out,
-    stride_logic_h, stride_logic_bs,
-    stride_prob_h, stride_prob_bs,
-    BLOCK_SIZE: tl.constexpr
+    stride_logic_h,
+    stride_logic_bs,
+    stride_prob_h,
+    stride_prob_bs,
+    BLOCK_SIZE: tl.constexpr,
 ):
     cur_batch = tl.program_id(0)
     cur_head = tl.program_id(1)
@@ -19,16 +23,22 @@ def _fwd_kernel_token_softmax(
     cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
     cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
 
-    row = tl.load(Logics + cur_head * stride_logic_h + (cur_batch_in_all_start_index + col_offsets) * stride_logic_bs,
-                  mask=col_offsets < cur_batch_seq_len, other=-float('inf')).to(tl.float32)
+    row = tl.load(
+        Logics + cur_head * stride_logic_h + (cur_batch_in_all_start_index + col_offsets) * stride_logic_bs,
+        mask=col_offsets < cur_batch_seq_len,
+        other=-float("inf"),
+    ).to(tl.float32)
 
     row_minus_max = row - tl.max(row, axis=0)
     numerator = tl.exp(row_minus_max)
     denominator = tl.sum(numerator, axis=0)
     softmax_output = numerator / denominator
 
-    tl.store(Prob_Out + cur_head * stride_prob_h + (cur_batch_in_all_start_index + col_offsets)
-             * stride_prob_bs, softmax_output, mask=col_offsets < cur_batch_seq_len)
+    tl.store(
+        Prob_Out + cur_head * stride_prob_h + (cur_batch_in_all_start_index + col_offsets) * stride_prob_bs,
+        softmax_output,
+        mask=col_offsets < cur_batch_seq_len,
+    )
     return
 
 
@@ -44,10 +54,14 @@ def token_softmax_fwd(Logics, B_Start_Loc, B_Seqlen, Prob_Out, max_input_len):
         num_warps = 16
 
     _fwd_kernel_token_softmax[(batch, head_num)](
-        Logics, B_Start_Loc, B_Seqlen,
+        Logics,
+        B_Start_Loc,
+        B_Seqlen,
         Prob_Out,
-        Logics.stride(0), Logics.stride(1),
-        Prob_Out.stride(0), Prob_Out.stride(1),
+        Logics.stride(0),
+        Logics.stride(1),
+        Prob_Out.stride(0),
+        Prob_Out.stride(1),
         num_warps=num_warps,
         BLOCK_SIZE=BLOCK_SIZE,
     )
@@ -59,7 +73,7 @@ def test1():
     import torch
 
     B, N_CTX, H, D = 4, 1025, 12, 128
-
+    del D
     dtype = torch.float16
 
     Logics = torch.empty((H, B * N_CTX), dtype=dtype, device="cuda").normal_(mean=0.1, std=10)
@@ -85,6 +99,7 @@ def test2():
     import torch
 
     B, N_CTX, H, D = 3, 1025, 12, 128
+    del D
 
     dtype = torch.float16
 
@@ -107,7 +122,7 @@ def test2():
     start = 0
     for i in range(B):
         end = start + b_seq_len[i]
-        torch_o = Logics[:, start: end].reshape(H * 1, -1).softmax(-1).reshape(H, 1 * b_seq_len[i])
+        torch_o = Logics[:, start:end].reshape(H * 1, -1).softmax(-1).reshape(H, 1 * b_seq_len[i])
         start = end
         torch_out.append(torch_o)
     torch_out = torch.cat(torch_out, dim=-1)
diff --git a/lightllm/models/bloom/triton_kernel/token_flashattention_nopad.py b/lightllm/common/basemodel/triton_kernel/alibi_att/token_flashattention_nopad.py
similarity index 100%
rename from lightllm/models/bloom/triton_kernel/token_flashattention_nopad.py
rename to lightllm/common/basemodel/triton_kernel/alibi_att/token_flashattention_nopad.py
diff --git a/lightllm/models/chatglm2/__init__.py b/lightllm/common/basemodel/triton_kernel/att/__init__.py
similarity index 100%
rename from lightllm/models/chatglm2/__init__.py
rename to lightllm/common/basemodel/triton_kernel/att/__init__.py
diff --git a/lightllm/models/chatglm2/layer_infer/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/__init__.py
similarity index 100%
rename from lightllm/models/chatglm2/layer_infer/__init__.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/__init__.py
diff --git a/lightllm/models/chatglm2/layer_weights/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/__init__.py
similarity index 100%
rename from lightllm/models/chatglm2/layer_weights/__init__.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/__init__.py
diff --git a/lightllm/models/chatglm2/triton_kernel/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/__init__.py
similarity index 100%
rename from lightllm/models/chatglm2/triton_kernel/__init__.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/__init__.py
diff --git a/lightllm/models/llama/triton_kernel/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py
similarity index 65%
rename from lightllm/models/llama/triton_kernel/gqa_flash_decoding.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py
index 67be7c968b..26ec3ebd71 100644
--- a/lightllm/models/llama/triton_kernel/gqa_flash_decoding.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py
@@ -2,11 +2,12 @@
 
 
 def gqa_token_decode_attention_flash_decoding(
-    q, infer_state, q_head_num, head_dim, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty
+    q: torch.Tensor, infer_state, cache_k: torch.Tensor, cache_v: torch.Tensor, out=None, alloc_tensor_func=torch.empty
 ):
     BLOCK_SEQ = 128
     batch_size = infer_state.batch_size
-    max_len_in_batch = infer_state.max_len_in_batch
+    max_kv_seq_len = infer_state.max_kv_seq_len
+    q_head_num, head_dim = q.shape[1], q.shape[2]
     calcu_shape1 = (batch_size, q_head_num, head_dim)
 
     from .gqa_flash_decoding_stage1 import flash_decode_stage1
@@ -15,10 +16,10 @@ def gqa_token_decode_attention_flash_decoding(
     o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out
 
     mid_o = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1, head_dim], dtype=torch.float32, device="cuda"
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1, head_dim], dtype=torch.float32, device="cuda"
     )
     mid_o_logexpsum = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1], dtype=torch.float32, device="cuda"
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1], dtype=torch.float32, device="cuda"
     )
 
     flash_decode_stage1(
@@ -28,7 +29,7 @@ def gqa_token_decode_attention_flash_decoding(
         infer_state.req_manager.req_to_token_indexs,
         infer_state.b_req_idx,
         infer_state.b_seq_len,
-        infer_state.max_len_in_batch,
+        infer_state.max_kv_seq_len,
         mid_o,
         mid_o_logexpsum,
         BLOCK_SEQ,
diff --git a/lightllm/models/llama/triton_kernel/gqa_flash_decoding_stage1.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py
similarity index 96%
rename from lightllm/models/llama/triton_kernel/gqa_flash_decoding_stage1.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py
index 320c2cf798..2814ff44bc 100644
--- a/lightllm/models/llama/triton_kernel/gqa_flash_decoding_stage1.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py
@@ -123,8 +123,18 @@ def _fwd_kernel_flash_decode_stage1(
 
 @torch.no_grad()
 def flash_decode_stage1(
-    q, k, v, Req_to_tokens, B_req_idx, B_Seqlen, max_len_in_batch, mid_out, mid_out_logsumexp, block_seq
+    q,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    Req_to_tokens,
+    B_req_idx,
+    B_Seqlen,
+    max_len_in_batch,
+    mid_out,
+    mid_out_logsumexp,
+    block_seq,
 ):
+    assert k.stride() == v.stride()
     BLOCK_SEQ = block_seq
     BLOCK_N = 16
     assert BLOCK_SEQ % BLOCK_N == 0
diff --git a/lightllm/models/llama/triton_kernel/gqa_flash_decoding_stage2.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage2.py
similarity index 65%
rename from lightllm/models/llama/triton_kernel/gqa_flash_decoding_stage2.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage2.py
index 81227f967b..101e99dde5 100644
--- a/lightllm/models/llama/triton_kernel/gqa_flash_decoding_stage2.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage2.py
@@ -6,14 +6,22 @@
 @triton.jit
 def _fwd_kernel_flash_decode_stage2(
     B_Seqlen,
-    Mid_O, # [batch, head, seq_block_num, head_dim]
-    Mid_O_LogExpSum, #[batch, head, seq_block_num]
-    O, #[batch, head, head_dim]
-    stride_mid_ob, stride_mid_oh, stride_mid_os, stride_mid_od,
-    stride_mid_o_eb, stride_mid_o_eh, stride_mid_o_es,
-    stride_obs, stride_oh, stride_od,
+    Mid_O,  # [batch, head, seq_block_num, head_dim]
+    Mid_O_LogExpSum,  # [batch, head, seq_block_num]
+    O,  # [batch, head, head_dim]
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_mid_od,
+    stride_mid_o_eb,
+    stride_mid_o_eh,
+    stride_mid_o_es,
+    stride_obs,
+    stride_oh,
+    stride_od,
     BLOCK_SEQ: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr):
+    BLOCK_DMODEL: tl.constexpr,
+):
     cur_batch = tl.program_id(0)
     cur_head = tl.program_id(1)
 
@@ -32,33 +40,43 @@ def _fwd_kernel_flash_decode_stage2(
         tv = tl.load(Mid_O + offs_v + block_seq_n * stride_mid_os)
         tlogic = tl.load(Mid_O_LogExpSum + offs_logic + block_seq_n)
         new_max_logic = tl.maximum(tlogic, max_logic)
-        
+
         old_scale = tl.exp(max_logic - new_max_logic)
         acc *= old_scale
         exp_logic = tl.exp(tlogic - new_max_logic)
         acc += exp_logic * tv
         sum_exp = sum_exp * old_scale + exp_logic
         max_logic = new_max_logic
-    
+
     tl.store(O + cur_batch * stride_obs + cur_head * stride_oh + offs_d, acc / sum_exp)
     return
 
 
 @torch.no_grad()
-def flash_decode_stage2(mid_out, mid_out_logexpsum, B_Seqlen, O, block_seq):
+def flash_decode_stage2(mid_out, mid_out_logexpsum, B_Seqlen, out, block_seq):
     Lk = mid_out.shape[-1]
     assert Lk in {16, 32, 64, 128}
     batch, head_num = mid_out.shape[0], mid_out.shape[1]
     grid = (batch, head_num)
-    
+
     _fwd_kernel_flash_decode_stage2[grid](
-        B_Seqlen, mid_out, mid_out_logexpsum, O,
-        mid_out.stride(0), mid_out.stride(1), mid_out.stride(2), mid_out.stride(3),
-        mid_out_logexpsum.stride(0), mid_out_logexpsum.stride(1), mid_out_logexpsum.stride(2),
-        O.stride(0), O.stride(1), O.stride(2),
+        B_Seqlen,
+        mid_out,
+        mid_out_logexpsum,
+        out,
+        mid_out.stride(0),
+        mid_out.stride(1),
+        mid_out.stride(2),
+        mid_out.stride(3),
+        mid_out_logexpsum.stride(0),
+        mid_out_logexpsum.stride(1),
+        mid_out_logexpsum.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
         BLOCK_SEQ=block_seq,
         BLOCK_DMODEL=Lk,
         num_warps=4,
         num_stages=2,
     )
-    return
\ No newline at end of file
+    return
diff --git a/lightllm/models/llama/triton_kernel/gqa_flash_decoding_vsm.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py
similarity index 99%
rename from lightllm/models/llama/triton_kernel/gqa_flash_decoding_vsm.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py
index 850d4185c3..6a9bb79c7d 100644
--- a/lightllm/models/llama/triton_kernel/gqa_flash_decoding_vsm.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py
@@ -421,7 +421,7 @@ def gqa_token_decode_attention_flash_decoding_vsm(
 
     if not run_config:
         if torch.cuda.is_current_stream_capturing():
-            avg_seq_len_in_batch = infer_state.max_len_in_batch
+            avg_seq_len_in_batch = infer_state.max_kv_seq_len
         else:
             avg_seq_len_in_batch = infer_state.total_token_num // batch_size
 
diff --git a/lightllm/models/llama/triton_kernel/gqa_decode_flashattention_nopad.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/gqa_decode_flashattention_nopad.py
similarity index 100%
rename from lightllm/models/llama/triton_kernel/gqa_decode_flashattention_nopad.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/gqa_decode_flashattention_nopad.py
diff --git a/lightllm/models/cohere/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/__init__.py
similarity index 100%
rename from lightllm/models/cohere/__init__.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/__init__.py
diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py
new file mode 100644
index 0000000000..212825a962
--- /dev/null
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py
@@ -0,0 +1,200 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def int4_to_float(k_int8, k_scale, offs_d):
+    k_int8 = k_int8.to(tl.uint8, bitcast=True)
+    k_high = (k_int8 & 0xF0) >> 4
+    k_low = k_int8 & 0x0F
+    k_high = k_high.to(tl.int8, bitcast=True)
+    k_low = k_low.to(tl.int8, bitcast=True)
+    k_high -= 7
+    k_low -= 7
+    k_int4 = tl.where(
+        offs_d[None, :] % 2 == 0,
+        k_low,
+        k_high,
+    )
+    k = k_int4.to(k_scale.dtype) * k_scale
+    return k
+
+
+@triton.jit
+def _fwd_kernel_flash_decode_stage1(
+    Q,
+    K,
+    K_scale,
+    V,
+    V_scale,
+    sm_scale,
+    Req_to_tokens,
+    B_req_idx,
+    B_Seqlen,
+    Mid_O,  # [batch, head, seq_block_num, head_dim]
+    Mid_O_LogExpSum,  # [batch, head, seq_block_num]
+    stride_req_to_tokens_b,
+    stride_req_to_tokens_s,
+    stride_qbs,
+    stride_qh,
+    stride_qd,
+    stride_kbs,
+    stride_kh,
+    stride_kd,
+    stride_vbs,
+    stride_vh,
+    stride_vd,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_mid_od,
+    stride_mid_o_eb,
+    stride_mid_o_eh,
+    stride_mid_o_es,
+    gqa_group_size,
+    quant_group_size,
+    BLOCK_SEQ: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    seq_start_block = tl.program_id(2)
+    cur_kv_head = cur_head // gqa_group_size
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
+    cur_batch_start_index = seq_start_block * BLOCK_SEQ
+    cur_batch_end_index = tl.minimum(cur_batch_seq_len, cur_batch_start_index + BLOCK_SEQ)
+
+    off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d
+
+    block_n_size = (
+        tl.where(
+            cur_batch_end_index - cur_batch_start_index <= 0,
+            0,
+            cur_batch_end_index - cur_batch_start_index + BLOCK_N - 1,
+        )
+        // BLOCK_N
+    )
+
+    offs_n = cur_batch_start_index + tl.arange(0, BLOCK_N)
+
+    q = tl.load(Q + off_q)
+
+    sum_exp = 0.0
+    max_logic = -float("inf")
+    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
+
+    for start_n in range(0, block_n_size, 1):
+        offs_n_new = start_n * BLOCK_N + offs_n
+        k_loc = tl.load(
+            Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_n_new,
+            mask=offs_n_new < cur_batch_end_index,
+            other=0,
+        )
+        k_loc = k_loc.to(tl.int64)
+        off_k = k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :] // 2
+        off_k_scale = off_k // (quant_group_size // 2)
+        k_int8 = tl.load(K + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0)
+        k_scale = tl.load(K_scale + off_k_scale, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0)
+        k = int4_to_float(k_int8, k_scale, offs_d)
+
+        att_value = tl.sum(q[None, :] * k, 1)
+        att_value *= sm_scale
+        att_value = tl.where((offs_n_new < cur_batch_end_index), att_value, float("-inf"))
+        v_int8 = tl.load(V + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0)
+        v_scale = tl.load(V_scale + off_k_scale, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0)
+        v = int4_to_float(v_int8, v_scale, offs_d)
+
+        cur_max_logic = tl.max(att_value, axis=0)
+        new_max_logic = tl.maximum(cur_max_logic, max_logic)
+
+        exp_logic = tl.exp(att_value - new_max_logic)
+        logic_scale = tl.exp(max_logic - new_max_logic)
+        acc *= logic_scale
+        acc += tl.sum(exp_logic[:, None] * v, axis=0)
+
+        sum_exp = sum_exp * logic_scale + tl.sum(exp_logic, axis=0)
+        max_logic = new_max_logic
+
+    need_store = tl.where(block_n_size == 0, 0, 1)
+    for _ in range(0, need_store, 1):
+        off_mid_o = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + seq_start_block * stride_mid_os + offs_d
+        off_mid_o_logexpsum = cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh + seq_start_block
+        tl.store(Mid_O + off_mid_o, acc / sum_exp)
+        tl.store(Mid_O_LogExpSum + off_mid_o_logexpsum, max_logic + tl.log(sum_exp))
+    return
+
+
+@torch.no_grad()
+def int4kv_flash_decode_stage1(
+    q,
+    k,
+    k_scale,
+    v,
+    v_scale,
+    Req_to_tokens,
+    B_req_idx,
+    B_Seqlen,
+    max_len_in_batch,
+    mid_out,
+    mid_out_logsumexp,
+    block_seq,
+):
+    BLOCK_SEQ = block_seq
+    BLOCK_N = 16
+    assert BLOCK_SEQ % BLOCK_N == 0
+    # shape constraints
+    Lq, Lk = q.shape[-1], k.shape[-1] * 2
+    assert Lq == Lk
+    assert Lk in {16, 32, 64, 128}
+    sm_scale = 1.0 / (Lk ** 0.5)
+    batch, head_num = B_req_idx.shape[0], q.shape[1]
+    grid = (batch, head_num, triton.cdiv(max_len_in_batch, BLOCK_SEQ))
+    gqa_group_size = q.shape[1] // k.shape[1]
+    quant_group_size = Lk // k_scale.shape[-1]
+    assert triton.next_power_of_2(quant_group_size) == quant_group_size
+    assert k.stride() == v.stride()
+    # TODO 优化为gqa使用tensor core的实现，速度更快。
+    _fwd_kernel_flash_decode_stage1[grid](
+        q,
+        k,
+        k_scale,
+        v,
+        v_scale,
+        sm_scale,
+        Req_to_tokens,
+        B_req_idx,
+        B_Seqlen,
+        mid_out,
+        mid_out_logsumexp,
+        Req_to_tokens.stride(0),
+        Req_to_tokens.stride(1),
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        v.stride(0),
+        v.stride(1),
+        v.stride(2),
+        mid_out.stride(0),
+        mid_out.stride(1),
+        mid_out.stride(2),
+        mid_out.stride(3),
+        mid_out_logsumexp.stride(0),
+        mid_out_logsumexp.stride(1),
+        mid_out_logsumexp.stride(2),
+        gqa_group_size,
+        quant_group_size,
+        BLOCK_SEQ=BLOCK_SEQ,
+        BLOCK_DMODEL=Lk,
+        BLOCK_N=BLOCK_N,
+        num_warps=4,
+        num_stages=2,
+    )
+    return
diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py
new file mode 100644
index 0000000000..a5a054b93a
--- /dev/null
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py
@@ -0,0 +1,50 @@
+import torch
+
+
+def token_decode_attention_flash_decoding(
+    q,
+    infer_state,
+    cache_k,
+    cache_k_scale,
+    cache_v,
+    cache_v_scale,
+    out=None,
+    alloc_tensor_func=torch.empty,
+):
+    BLOCK_SEQ = 256
+    batch_size = infer_state.batch_size
+    max_kv_seq_len = infer_state.max_kv_seq_len
+    q_head_num = q.shape[1]
+    head_dim = q.shape[2]
+    calcu_shape1 = (batch_size, q_head_num, head_dim)
+
+    from ..mha.flash_decoding.flash_decoding_stage2 import flash_decode_stage2
+
+    o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out
+
+    mid_o = alloc_tensor_func(
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1, head_dim], dtype=q.dtype, device="cuda"
+    )
+    mid_o_logexpsum = alloc_tensor_func(
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1], dtype=q.dtype, device="cuda"
+    )
+
+    from .int4kv_flash_decoding_stage1 import int4kv_flash_decode_stage1
+
+    int4kv_flash_decode_stage1(
+        q=q.view(calcu_shape1),
+        k=cache_k,
+        k_scale=cache_k_scale,
+        v=cache_v,
+        v_scale=cache_v_scale,
+        Req_to_tokens=infer_state.req_manager.req_to_token_indexs,
+        B_req_idx=infer_state.b_req_idx,
+        B_Seqlen=infer_state.b_seq_len,
+        max_len_in_batch=infer_state.max_kv_seq_len,
+        mid_out=mid_o,
+        mid_out_logsumexp=mid_o_logexpsum,
+        block_seq=BLOCK_SEQ,
+    )
+
+    flash_decode_stage2(mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ)
+    return o_tensor
diff --git a/lightllm/models/cohere/layer_infer/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/__init__.py
similarity index 100%
rename from lightllm/models/cohere/layer_infer/__init__.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/__init__.py
diff --git a/lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding_diverse.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py
similarity index 67%
rename from lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding_diverse.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py
index 84054bf867..ad6a8b5b3a 100644
--- a/lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding_diverse.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py
@@ -2,16 +2,15 @@
 import torch
 from lightllm.utils.light_utils import HAS_LIGHTLLM_KERNEL, light_ops
 from lightllm.common.basemodel.infer_struct import InferStateInfo
-from .ppl_int8kv_flash_decoding_diverse_stage1 import flash_decode_stage1
-from .ppl_int8kv_flash_decoding_diverse_stage3 import flash_diverse_decode_stage3
+from .int8kv_flash_decoding_diverse_stage1 import flash_decode_stage1
+from .int8kv_flash_decoding_diverse_stage2 import flash_decode_stage2
+from .int8kv_flash_decoding_diverse_stage3 import flash_diverse_decode_stage3
 from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size
 
 
 def token_decode_attention_flash_decoding(
     q,
     infer_state: InferStateInfo,
-    q_head_num,
-    head_dim,
     cache_k,
     cache_k_scale,
     cache_v,
@@ -28,18 +27,21 @@ def token_decode_attention_flash_decoding(
     stream1 = shared_streams_dict["stream1"]
     stream2 = shared_streams_dict["stream2"]
 
+    q_head_num = q.shape[1]
+    head_dim = q.shape[2]
+
     BLOCK_SEQ = 256
     batch_size = infer_state.batch_size
-    max_len_in_batch = infer_state.max_len_in_batch
+    max_kv_seq_len = infer_state.max_kv_seq_len
     calcu_shape1 = (batch_size, q_head_num, head_dim)
 
     o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out
 
     mid_o = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 2, head_dim], dtype=q.dtype, device="cuda"
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 2, head_dim], dtype=torch.float32, device="cuda"
     )
     mid_o_logexpsum = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 2], dtype=q.dtype, device="cuda"
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 2], dtype=torch.float32, device="cuda"
     )
 
     current_stream = torch.cuda.current_stream()
@@ -56,7 +58,7 @@ def token_decode_attention_flash_decoding(
             B_req_idx=infer_state.b_req_idx,
             b_shared_seq_len=infer_state.b_shared_seq_len,
             b_mark_shared_group=infer_state.b_mark_shared_group,
-            max_len_in_batch=infer_state.max_len_in_batch,
+            max_len_in_batch=infer_state.max_kv_seq_len,
             mid_out=mid_o,
             mid_out_logsumexp=mid_o_logexpsum,
             block_seq=BLOCK_SEQ,
@@ -64,21 +66,20 @@ def token_decode_attention_flash_decoding(
         )
     stream2.wait_stream(current_stream)
     with torch.cuda.stream(stream2):
-        light_ops.group8_int8kv_flashdecoding_diverse_stage2(
-            BLOCK_SEQ,
-            mid_o,
-            mid_o_logexpsum,
-            1.0 / (head_dim ** 0.5),
-            q.view(calcu_shape1),
-            cache_k,
-            cache_k_scale,
-            cache_v,
-            cache_v_scale,
-            infer_state.req_manager.req_to_token_indexs,
-            infer_state.b_req_idx,
-            infer_state.b_seq_len,
-            infer_state.b_shared_seq_len,
-            infer_state.max_len_in_batch,
+        flash_decode_stage2(
+            q=q.view(calcu_shape1),
+            k=cache_k,
+            k_scale=cache_k_scale,
+            v=cache_v,
+            v_scale=cache_v_scale,
+            Req_to_tokens=infer_state.req_manager.req_to_token_indexs,
+            B_req_idx=infer_state.b_req_idx,
+            B_Seqlen=infer_state.b_seq_len,
+            b_shared_seq_len=infer_state.b_shared_seq_len,
+            max_len_in_batch=infer_state.max_kv_seq_len,
+            mid_out=mid_o,
+            mid_out_logsumexp=mid_o_logexpsum,
+            block_seq=BLOCK_SEQ,
         )
 
     current_stream.wait_stream(stream1)
diff --git a/lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding_diverse_stage1.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py
similarity index 89%
rename from lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding_diverse_stage1.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py
index 8b3423ce99..295ae66ab3 100644
--- a/lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding_diverse_stage1.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py
@@ -9,7 +9,7 @@
 
 
 class GQADiverseDecodeStage1KernelConfig(KernelConfigs):
-    kernel_name: str = "_fwd_kernel_flash_decode_diverse_stage1:v1"
+    kernel_name: str = "_fwd_kernel_flash_decode_diverse_stage1:v2"
 
     @classmethod
     @lru_cache(maxsize=200)
@@ -113,6 +113,7 @@ def _fwd_kernel_flash_decode_diverse_stage1(
     BLOCK_N: tl.constexpr,
     BLOCK_BATCH: tl.constexpr,
     KV_QUANT_GROUP_SIZE: tl.constexpr,
+    NUM_GROUPS: tl.constexpr,
 ):
     cur_batch = tl.program_id(0)
     shared_batch_group_size = tl.load(b_mark_shared_group + cur_batch)
@@ -128,6 +129,7 @@ def _fwd_kernel_flash_decode_diverse_stage1(
     cur_q_head_range = tl.where(cur_q_head_range < q_head_end_index, cur_q_head_range, cur_kv_head * gqa_group_size)
 
     offs_d = tl.arange(0, BLOCK_HEADDIM)
+    offs_d_scale = tl.arange(0, NUM_GROUPS)
     cur_batch_seq_len = tl.load(b_shared_seq_len + cur_batch)
     cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
     cur_batch_start_index = seq_start_block * BLOCK_SEQ
@@ -162,25 +164,37 @@ def _fwd_kernel_flash_decode_diverse_stage1(
             mask=n_mask,
             other=0,
         ).to(tl.int64)
-        off_k = k_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None]
-        off_k_scale = off_k // KV_QUANT_GROUP_SIZE
+        off_k_base = k_loc * stride_kbs + cur_kv_head * stride_kh
+        # (128, 16)
+        off_k = off_k_base[None, :] + offs_d[:, None]
+        # off_k_scale = off_k // KV_QUANT_GROUP_SIZE
+        # (16, 16)
+        off_k_scale = off_k_base[None, :] // KV_QUANT_GROUP_SIZE + offs_d_scale[:, None]
         k = tl.load(K + off_k, mask=n_mask[None, :], other=0)
+        k = tl.reshape(k, (NUM_GROUPS, KV_QUANT_GROUP_SIZE, BLOCK_N))
         k_scale = tl.load(K_scale + off_k_scale, mask=n_mask[None, :], other=0.0)
+        k_scale = tl.reshape(k_scale, (NUM_GROUPS, 1, BLOCK_N))
         k = k * k_scale
+        k = tl.reshape(k, (BLOCK_HEADDIM, BLOCK_N))
         att_value = tl.dot(q, k.to(q.dtype))
         att_value *= sm_scale
         att_value = tl.where(n_mask[None, :], att_value, float("-inf"))
+        off_v = k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :]
         v = tl.load(
-            V + off_k.T,
+            V + off_v,
             mask=n_mask[:, None],
             other=0,
         )
+        v = tl.reshape(v, (BLOCK_N, NUM_GROUPS, KV_QUANT_GROUP_SIZE))
         v_scale = tl.load(
-            V_scale + off_k_scale.T,
-            mask=n_mask[:, None],
+            V_scale + off_k_scale,
+            mask=n_mask[None, :],
             other=0.0,
         )
+        v_scale = tl.trans(v_scale)
+        v_scale = tl.reshape(v_scale, (BLOCK_N, NUM_GROUPS, 1))
         v = v * v_scale
+        v = tl.reshape(v, (BLOCK_N, BLOCK_HEADDIM))
 
         cur_max_logic = tl.max(att_value, axis=1)
         new_max_logic = tl.maximum(cur_max_logic, max_logic)
@@ -269,12 +283,16 @@ def flash_decode_stage1(
     gqa_group_size = q.shape[1] // k.shape[1]
     assert triton.next_power_of_2(Lk) == Lk
     KV_QUANT_GROUP_SIZE = v.shape[-1] // v_scale.shape[-1]
-    assert KV_QUANT_GROUP_SIZE == 8
+    assert triton.next_power_of_2(KV_QUANT_GROUP_SIZE) == KV_QUANT_GROUP_SIZE
     BLOCK_HEAD = triton.next_power_of_2(gqa_group_size)
     BLOCK_BATCH = triton.next_power_of_2(max_batch_group_size)
     if BLOCK_HEAD * BLOCK_BATCH < 16:
         BLOCK_BATCH = 16 // BLOCK_HEAD
+    assert k.stride() == v.stride()
+    NUM_GROUPS = Lk // KV_QUANT_GROUP_SIZE
+    assert triton.next_power_of_2(NUM_GROUPS) == NUM_GROUPS
 
+    assert k.stride() == v.stride()
     _fwd_kernel_flash_decode_diverse_stage1[grid](
         Q=q,
         stride_qbs=q.stride(0),
@@ -313,6 +331,7 @@ def flash_decode_stage1(
         BLOCK_N=BLOCK_N,
         BLOCK_BATCH=BLOCK_BATCH,
         KV_QUANT_GROUP_SIZE=KV_QUANT_GROUP_SIZE,
+        NUM_GROUPS=NUM_GROUPS,
         num_warps=num_warps,
         num_stages=num_stages,
     )
diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py
new file mode 100644
index 0000000000..f5c0b9c395
--- /dev/null
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py
@@ -0,0 +1,306 @@
+import torch
+import triton
+import triton.language as tl
+from typing import Optional
+
+from lightllm.common.kernel_config import KernelConfigs
+from frozendict import frozendict
+from functools import lru_cache
+from typing import Dict
+from lightllm.common.triton_utils.autotuner import autotune, Autotuner
+
+
+class GQADiverseDecodeStage2KernelConfig(KernelConfigs):
+    kernel_name: str = "_fwd_kernel_flash_decode_diverse_stage2:v1"
+
+    @classmethod
+    @lru_cache(maxsize=200)
+    def try_to_get_best_config(
+        cls,
+        batch_size: int,
+        avg_seq_len_in_batch: int,
+        gqa_group_size: int,
+        q_head_dim: int,
+        block_seq: int,
+        out_dtype: str,
+    ) -> dict:
+        key_params = {
+            "gqa_group_size": gqa_group_size,
+            "q_head_dim": q_head_dim,
+            "block_seq": block_seq,
+            "out_dtype": str(out_dtype),
+        }
+        key_params = frozendict(key_params)
+
+        finded_config = cls.get_the_config(key_params)
+
+        if finded_config:
+            batch_size_config: dict = finded_config[
+                min(
+                    finded_config.keys(),
+                    key=lambda x: abs(int(x) - avg_seq_len_in_batch),
+                )
+            ]
+            config = batch_size_config[min(batch_size_config.keys(), key=lambda x: abs(int(x) - batch_size))]
+
+            return config
+        else:
+            config = {
+                "BLOCK_N": 16,
+                "num_warps": 2,
+                "num_stages": 2,
+            }
+        return config
+
+    @classmethod
+    def save_config(
+        cls,
+        gqa_group_size: int,
+        q_head_dim: int,
+        block_seq: int,
+        out_dtype: str,
+        config_json: Dict[int, Dict[int, Dict]],
+    ):
+        key_params = {
+            "gqa_group_size": gqa_group_size,
+            "q_head_dim": q_head_dim,
+            "block_seq": block_seq,
+            "out_dtype": str(out_dtype),
+        }
+        key_params = frozendict(key_params)
+
+        return cls.store_config(key_params, config_json)
+
+
+@triton.jit
+def _fwd_kernel_flash_decode_diverse_stage2(
+    Q,
+    stride_qbs,
+    stride_qh,
+    stride_qd,
+    K,
+    K_scale,
+    stride_kbs,
+    stride_kh,
+    stride_kd,
+    V,
+    V_scale,
+    stride_vbs,
+    stride_vh,
+    stride_vd,
+    sm_scale,
+    Req_to_tokens,
+    stride_req_to_tokens_b,
+    stride_req_to_tokens_s,
+    B_req_idx,
+    B_Seqlen,
+    b_shared_seq_len,
+    Mid_O,  # [batch, head, seq_block_num, head_dim]
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_mid_od,
+    Mid_O_LogExpSum,  # [batch, head, seq_block_num]
+    stride_mid_o_eb,
+    stride_mid_o_eh,
+    stride_mid_o_es,
+    gqa_group_size: tl.constexpr,
+    BLOCK_SEQ: tl.constexpr,
+    BLOCK_HEADDIM: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    KV_QUANT_GROUP_SIZE: tl.constexpr,
+    NUM_GROUPS: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_kv_head = tl.program_id(1)
+    seq_start_block = tl.program_id(2)
+
+    cur_q_head_range = cur_kv_head * gqa_group_size + tl.arange(0, gqa_group_size)
+
+    offs_d = tl.arange(0, BLOCK_HEADDIM)
+    offs_d_scale = tl.arange(0, NUM_GROUPS)
+    cur_batch_shared_len = tl.load(b_shared_seq_len + cur_batch)
+    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
+    cur_batch_start_index = seq_start_block * BLOCK_SEQ + cur_batch_shared_len
+    cur_batch_end_index = tl.minimum(cur_batch_seq_len, cur_batch_start_index + BLOCK_SEQ)
+    store_seq_block = seq_start_block + tl.cdiv(cur_batch_shared_len, BLOCK_SEQ)
+
+    off_q = cur_batch * stride_qbs + cur_q_head_range[:, None] * stride_qh + offs_d[None, :]
+
+    block_n_size = tl.cdiv(
+        tl.where(cur_batch_end_index - cur_batch_start_index <= 0, 0, cur_batch_end_index - cur_batch_start_index),
+        BLOCK_N,
+    )
+
+    if block_n_size == 0:
+        return
+
+    offs_n = cur_batch_start_index + tl.arange(0, BLOCK_N)
+    q = tl.load(Q + off_q)
+
+    sum_exp = tl.zeros([gqa_group_size], dtype=tl.float32)
+    max_logic = tl.zeros([gqa_group_size], dtype=tl.float32) - float("inf")
+    acc = tl.zeros([gqa_group_size, BLOCK_HEADDIM], dtype=tl.float32)
+
+    for start_n in range(0, block_n_size, 1):
+        offs_n_new = start_n * BLOCK_N + offs_n
+        n_mask = offs_n_new < cur_batch_end_index
+        k_loc = tl.load(
+            Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_n_new,
+            mask=n_mask,
+            other=0,
+        ).to(tl.int64)
+        off_k_base = k_loc * stride_kbs + cur_kv_head * stride_kh
+        # (128, 16)
+        off_k = off_k_base[None, :] + offs_d[:, None]
+        # off_k_scale = off_k // KV_QUANT_GROUP_SIZE
+        # (16, 16)
+        off_k_scale = off_k_base[None, :] // KV_QUANT_GROUP_SIZE + offs_d_scale[:, None]
+        k = tl.load(K + off_k, mask=n_mask[None, :], other=0)
+        k = tl.reshape(k, (NUM_GROUPS, KV_QUANT_GROUP_SIZE, BLOCK_N))
+        k_scale = tl.load(K_scale + off_k_scale, mask=n_mask[None, :], other=0.0)
+        k_scale = tl.reshape(k_scale, (NUM_GROUPS, 1, BLOCK_N))
+        k = k * k_scale
+        k = tl.reshape(k, (BLOCK_HEADDIM, BLOCK_N))
+        # q (4, 128) k (128, BLOCK_N)
+        att_value = tl.dot(q, k.to(q.dtype))
+        att_value *= sm_scale
+        att_value = tl.where(n_mask[None, :], att_value, float("-inf"))
+        off_v = k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :]
+        v = tl.load(
+            V + off_v,
+            mask=n_mask[:, None],
+            other=0,
+        )
+        v = tl.reshape(v, (BLOCK_N, NUM_GROUPS, KV_QUANT_GROUP_SIZE))
+        v_scale = tl.load(
+            V_scale + off_k_scale,
+            mask=n_mask[None, :],
+            other=0.0,
+        )
+        v_scale = tl.trans(v_scale)
+        v_scale = tl.reshape(v_scale, (BLOCK_N, NUM_GROUPS, 1))
+        v = v * v_scale
+        v = tl.reshape(v, (BLOCK_N, BLOCK_HEADDIM))
+
+        cur_max_logic = tl.max(att_value, axis=1)
+        new_max_logic = tl.maximum(cur_max_logic, max_logic)
+
+        exp_logic = tl.exp(att_value - new_max_logic[:, None])
+        logic_scale = tl.exp(max_logic - new_max_logic)
+        acc *= logic_scale[:, None]
+        acc += tl.dot(exp_logic.to(q.dtype), v.to(q.dtype))
+
+        sum_exp = sum_exp * logic_scale + tl.sum(exp_logic, axis=1)
+        max_logic = new_max_logic
+
+    off_mid_o = (
+        cur_batch * stride_mid_ob
+        + cur_q_head_range[:, None] * stride_mid_oh
+        + store_seq_block * stride_mid_os
+        + offs_d[None, :]
+    )
+    off_mid_o_logexpsum = cur_batch * stride_mid_o_eb + cur_q_head_range * stride_mid_o_eh + store_seq_block
+    tl.store(
+        Mid_O + off_mid_o,
+        (acc / sum_exp[:, None]),
+    )
+    tl.store(
+        Mid_O_LogExpSum + off_mid_o_logexpsum,
+        (max_logic + tl.log(sum_exp)),
+    )
+    return
+
+
+@torch.no_grad()
+def flash_decode_stage2(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    k_scale: torch.Tensor,
+    v: torch.Tensor,
+    v_scale: torch.Tensor,
+    Req_to_tokens: torch.Tensor,
+    B_req_idx: torch.Tensor,
+    B_Seqlen: torch.Tensor,
+    b_shared_seq_len: torch.Tensor,
+    max_len_in_batch: int,
+    mid_out: torch.Tensor,
+    mid_out_logsumexp: torch.Tensor,
+    block_seq: int,
+    run_config: Optional[dict] = None,
+):
+    if not run_config:
+        run_config = GQADiverseDecodeStage2KernelConfig.try_to_get_best_config(
+            batch_size=int(q.shape[0]),
+            avg_seq_len_in_batch=max_len_in_batch,
+            gqa_group_size=int(q.shape[1] // k.shape[1]),
+            q_head_dim=int(q.shape[2]),
+            block_seq=block_seq,
+            out_dtype=q.dtype,
+        )
+
+    BLOCK_N = run_config["BLOCK_N"]
+    num_warps = run_config["num_warps"]
+    num_stages = run_config["num_stages"]
+
+    assert q.dim() == 3 and k.dim() == 3 and v.dim() == 3
+    BLOCK_SEQ = block_seq
+    assert BLOCK_SEQ % BLOCK_N == 0
+    # shape constraints
+    Lq, Lk = q.shape[-1], k.shape[-1]
+    assert Lq == Lk
+    assert Lk in {16, 32, 64, 128}
+    sm_scale = 1.0 / (Lk ** 0.5)
+    batch, kv_head_num = B_req_idx.shape[0], k.shape[1]
+    grid = (batch, kv_head_num, triton.cdiv(max_len_in_batch, BLOCK_SEQ))
+    gqa_group_size = q.shape[1] // k.shape[1]
+    assert triton.next_power_of_2(Lk) == Lk
+    KV_QUANT_GROUP_SIZE = v.shape[-1] // v_scale.shape[-1]
+    assert KV_QUANT_GROUP_SIZE == 8
+    NUM_GROUPS = Lk // KV_QUANT_GROUP_SIZE
+    assert triton.next_power_of_2(NUM_GROUPS) == NUM_GROUPS
+
+    assert k.stride() == v.stride()
+
+    _fwd_kernel_flash_decode_diverse_stage2[grid](
+        Q=q,
+        stride_qbs=q.stride(0),
+        stride_qh=q.stride(1),
+        stride_qd=q.stride(2),
+        K=k,
+        K_scale=k_scale,
+        stride_kbs=k.stride(0),
+        stride_kh=k.stride(1),
+        stride_kd=k.stride(2),
+        V=v,
+        V_scale=v_scale,
+        stride_vbs=v.stride(0),
+        stride_vh=v.stride(1),
+        stride_vd=v.stride(2),
+        sm_scale=sm_scale,
+        Req_to_tokens=Req_to_tokens,
+        stride_req_to_tokens_b=Req_to_tokens.stride(0),
+        stride_req_to_tokens_s=Req_to_tokens.stride(1),
+        B_req_idx=B_req_idx,
+        B_Seqlen=B_Seqlen,
+        b_shared_seq_len=b_shared_seq_len,
+        Mid_O=mid_out,
+        stride_mid_ob=mid_out.stride(0),
+        stride_mid_oh=mid_out.stride(1),
+        stride_mid_os=mid_out.stride(2),
+        stride_mid_od=mid_out.stride(3),
+        Mid_O_LogExpSum=mid_out_logsumexp,  # [batch, head, seq_block_num]
+        stride_mid_o_eb=mid_out_logsumexp.stride(0),
+        stride_mid_o_eh=mid_out_logsumexp.stride(1),
+        stride_mid_o_es=mid_out_logsumexp.stride(2),
+        gqa_group_size=gqa_group_size,
+        BLOCK_SEQ=block_seq,
+        BLOCK_HEADDIM=Lk,
+        BLOCK_N=BLOCK_N,
+        KV_QUANT_GROUP_SIZE=KV_QUANT_GROUP_SIZE,
+        NUM_GROUPS=NUM_GROUPS,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    return
diff --git a/lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding_diverse_stage3.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage3.py
similarity index 100%
rename from lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding_diverse_stage3.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage3.py
diff --git a/lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/ppl_int8kv_flash_decoding.py
similarity index 71%
rename from lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/ppl_int8kv_flash_decoding.py
index 88e39b82fc..f51d611661 100644
--- a/lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/ppl_int8kv_flash_decoding.py
@@ -5,8 +5,6 @@
 def token_decode_attention_flash_decoding(
     q,
     infer_state,
-    q_head_num,
-    head_dim,
     cache_k,
     cache_k_scale,
     cache_v,
@@ -15,19 +13,20 @@ def token_decode_attention_flash_decoding(
     alloc_tensor_func=torch.empty,
 ):
     BLOCK_SEQ = 256
+    q_head_num, head_dim = q.shape[1], q.shape[2]
     batch_size = infer_state.batch_size
-    max_len_in_batch = infer_state.max_len_in_batch
+    max_kv_seq_len = infer_state.max_kv_seq_len
     calcu_shape1 = (batch_size, q_head_num, head_dim)
 
-    from .flash_decoding_stage2 import flash_decode_stage2
+    from ..mha.flash_decoding.flash_decoding_stage2 import flash_decode_stage2
 
     o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out
 
     mid_o = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1, head_dim], dtype=q.dtype, device="cuda"
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1, head_dim], dtype=q.dtype, device="cuda"
     )
     mid_o_logexpsum = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1], dtype=q.dtype, device="cuda"
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1], dtype=q.dtype, device="cuda"
     )
 
     light_ops.group8_int8kv_flashdecoding_stage1(
@@ -43,7 +42,7 @@ def token_decode_attention_flash_decoding(
         infer_state.req_manager.req_to_token_indexs,
         infer_state.b_req_idx,
         infer_state.b_seq_len,
-        infer_state.max_len_in_batch,
+        infer_state.max_kv_seq_len,
     )
 
     flash_decode_stage2(mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ)
diff --git a/lightllm/models/cohere/layer_weights/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/__init__.py
similarity index 100%
rename from lightllm/models/cohere/layer_weights/__init__.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/mha/__init__.py
diff --git a/lightllm/models/cohere/triton_kernels/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/__init__.py
similarity index 100%
rename from lightllm/models/cohere/triton_kernels/__init__.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/__init__.py
diff --git a/lightllm/models/phi3/triton_kernel/flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py
similarity index 63%
rename from lightllm/models/phi3/triton_kernel/flash_decoding.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py
index e47e308864..6c50fc3927 100644
--- a/lightllm/models/phi3/triton_kernel/flash_decoding.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py
@@ -1,12 +1,11 @@
 import torch
 
 
-def token_decode_attention_flash_decoding(
-    q, infer_state, q_head_num, head_dim, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty
-):
+def token_decode_attention_flash_decoding(q, infer_state, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty):
     BLOCK_SEQ = 256
     batch_size = infer_state.batch_size
-    max_len_in_batch = infer_state.max_len_in_batch
+    max_kv_seq_len = infer_state.max_kv_seq_len
+    q_head_num, head_dim = q.shape[1], q.shape[2]
     calcu_shape1 = (batch_size, q_head_num, head_dim)
 
     from .flash_decoding_stage1 import flash_decode_stage1
@@ -15,10 +14,10 @@ def token_decode_attention_flash_decoding(
     o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out
 
     mid_o = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1, head_dim], dtype=torch.float32, device="cuda"
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1, head_dim], dtype=torch.float32, device="cuda"
     )
     mid_o_logexpsum = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1], dtype=torch.float32, device="cuda"
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1], dtype=torch.float32, device="cuda"
     )
 
     flash_decode_stage1(
@@ -28,7 +27,7 @@ def token_decode_attention_flash_decoding(
         infer_state.req_manager.req_to_token_indexs,
         infer_state.b_req_idx,
         infer_state.b_seq_len,
-        infer_state.max_len_in_batch,
+        infer_state.max_kv_seq_len,
         mid_o,
         mid_o_logexpsum,
         BLOCK_SEQ,
diff --git a/lightllm/models/phi3/triton_kernel/flash_decoding_stage1.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage1.py
similarity index 88%
rename from lightllm/models/phi3/triton_kernel/flash_decoding_stage1.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage1.py
index f6d8b5abee..f41a5c8fde 100644
--- a/lightllm/models/phi3/triton_kernel/flash_decoding_stage1.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage1.py
@@ -33,7 +33,6 @@ def _fwd_kernel_flash_decode_stage1(
     stride_mid_o_eh,
     stride_mid_o_es,
     gqa_group_size,
-    head_dim,
     BLOCK_SEQ: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
@@ -62,7 +61,7 @@ def _fwd_kernel_flash_decode_stage1(
 
     offs_n = cur_batch_start_index + tl.arange(0, BLOCK_N)
 
-    q = tl.load(Q + off_q, mask=offs_d < head_dim, other=0.0)
+    q = tl.load(Q + off_q)
 
     sum_exp = 0.0
     max_logic = -float("inf")
@@ -75,16 +74,13 @@ def _fwd_kernel_flash_decode_stage1(
             mask=offs_n_new < cur_batch_end_index,
             other=0,
         )
+        k_loc = k_loc.to(tl.int64)
         off_k = k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :]
-        k = tl.load(
-            K + off_k, mask=(offs_n_new[:, None] < cur_batch_end_index) & (offs_d[None, :] < head_dim), other=0.0
-        )
+        k = tl.load(K + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0)
         att_value = tl.sum(q[None, :] * k, 1)
         att_value *= sm_scale
         att_value = tl.where(offs_n_new < cur_batch_end_index, att_value, float("-inf"))
-        v = tl.load(
-            V + off_k, mask=(offs_n_new[:, None] < cur_batch_end_index) & (offs_d[None, :] < head_dim), other=0.0
-        )
+        v = tl.load(V + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0)
 
         cur_max_logic = tl.max(att_value, axis=0)
         new_max_logic = tl.maximum(cur_max_logic, max_logic)
@@ -101,7 +97,7 @@ def _fwd_kernel_flash_decode_stage1(
     for _ in range(0, need_store, 1):
         off_mid_o = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + seq_start_block * stride_mid_os + offs_d
         off_mid_o_logexpsum = cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh + seq_start_block
-        tl.store(Mid_O + off_mid_o, acc / sum_exp, mask=offs_d < head_dim)
+        tl.store(Mid_O + off_mid_o, acc / sum_exp)
         tl.store(Mid_O_LogExpSum + off_mid_o_logexpsum, max_logic + tl.log(sum_exp))
     return
 
@@ -116,13 +112,12 @@ def flash_decode_stage1(
     # shape constraints
     Lq, Lk = q.shape[-1], k.shape[-1]
     assert Lq == Lk
-    head_dim = Lq
-    BLOCK_DMODEL = triton.next_power_of_2(head_dim)
+    assert Lk in {16, 32, 64, 128}
     sm_scale = 1.0 / (Lk ** 0.5)
     batch, head_num = B_req_idx.shape[0], q.shape[1]
     grid = (batch, head_num, triton.cdiv(max_len_in_batch, BLOCK_SEQ))
     gqa_group_size = q.shape[1] // k.shape[1]
-
+    assert k.stride() == v.stride()
     _fwd_kernel_flash_decode_stage1[grid](
         q,
         k,
@@ -152,9 +147,8 @@ def flash_decode_stage1(
         mid_out_logsumexp.stride(1),
         mid_out_logsumexp.stride(2),
         gqa_group_size,
-        head_dim,
         BLOCK_SEQ=BLOCK_SEQ,
-        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DMODEL=Lk,
         BLOCK_N=BLOCK_N,
         num_warps=1,
         num_stages=2,
diff --git a/lightllm/models/phi3/triton_kernel/flash_decoding_stage2.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage2.py
similarity index 81%
rename from lightllm/models/phi3/triton_kernel/flash_decoding_stage2.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage2.py
index a06ee54545..101e99dde5 100644
--- a/lightllm/models/phi3/triton_kernel/flash_decoding_stage2.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage2.py
@@ -8,7 +8,7 @@ def _fwd_kernel_flash_decode_stage2(
     B_Seqlen,
     Mid_O,  # [batch, head, seq_block_num, head_dim]
     Mid_O_LogExpSum,  # [batch, head, seq_block_num]
-    Out,  # [batch, head, head_dim]
+    O,  # [batch, head, head_dim]
     stride_mid_ob,
     stride_mid_oh,
     stride_mid_os,
@@ -19,7 +19,6 @@ def _fwd_kernel_flash_decode_stage2(
     stride_obs,
     stride_oh,
     stride_od,
-    head_dim,
     BLOCK_SEQ: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
 ):
@@ -38,7 +37,7 @@ def _fwd_kernel_flash_decode_stage2(
     offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d
     offs_logic = cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh
     for block_seq_n in range(0, block_n_size, 1):
-        tv = tl.load(Mid_O + offs_v + block_seq_n * stride_mid_os, mask=offs_d < head_dim, other=0.0)
+        tv = tl.load(Mid_O + offs_v + block_seq_n * stride_mid_os)
         tlogic = tl.load(Mid_O_LogExpSum + offs_logic + block_seq_n)
         new_max_logic = tl.maximum(tlogic, max_logic)
 
@@ -49,23 +48,22 @@ def _fwd_kernel_flash_decode_stage2(
         sum_exp = sum_exp * old_scale + exp_logic
         max_logic = new_max_logic
 
-    tl.store(Out + cur_batch * stride_obs + cur_head * stride_oh + offs_d, acc / sum_exp, mask=offs_d < head_dim)
+    tl.store(O + cur_batch * stride_obs + cur_head * stride_oh + offs_d, acc / sum_exp)
     return
 
 
 @torch.no_grad()
-def flash_decode_stage2(mid_out, mid_out_logexpsum, B_Seqlen, Out, block_seq):
+def flash_decode_stage2(mid_out, mid_out_logexpsum, B_Seqlen, out, block_seq):
     Lk = mid_out.shape[-1]
-    head_dim = Lk
+    assert Lk in {16, 32, 64, 128}
     batch, head_num = mid_out.shape[0], mid_out.shape[1]
-    BLOCK_DMODEL = triton.next_power_of_2(head_dim)
     grid = (batch, head_num)
 
     _fwd_kernel_flash_decode_stage2[grid](
         B_Seqlen,
         mid_out,
         mid_out_logexpsum,
-        Out,
+        out,
         mid_out.stride(0),
         mid_out.stride(1),
         mid_out.stride(2),
@@ -73,12 +71,11 @@ def flash_decode_stage2(mid_out, mid_out_logexpsum, B_Seqlen, Out, block_seq):
         mid_out_logexpsum.stride(0),
         mid_out_logexpsum.stride(1),
         mid_out_logexpsum.stride(2),
-        Out.stride(0),
-        Out.stride(1),
-        Out.stride(2),
-        head_dim,
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
         BLOCK_SEQ=block_seq,
-        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DMODEL=Lk,
         num_warps=4,
         num_stages=2,
     )
diff --git a/lightllm/models/mistral/triton_kernel/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/__init__.py
similarity index 100%
rename from lightllm/models/mistral/triton_kernel/__init__.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/__init__.py
diff --git a/lightllm/models/mistral/triton_kernel/token_attention_nopad_att1.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py
similarity index 63%
rename from lightllm/models/mistral/triton_kernel/token_attention_nopad_att1.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py
index 9a8261132a..9de2b82057 100644
--- a/lightllm/models/mistral/triton_kernel/token_attention_nopad_att1.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py
@@ -14,8 +14,6 @@ def _fwd_kernel_token_att1(
     B_req_idx,
     B_Start_Loc,
     B_Seqlen,
-    B_Att_Start_Loc,
-    B_Att_Seqlen,
     Att_Out,
     stride_req_to_tokens_b,
     stride_req_to_tokens_s,
@@ -28,7 +26,6 @@ def _fwd_kernel_token_att1(
     att_stride_h,
     att_stride_bs,
     kv_group_num,
-    sliding_window,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
 ):
@@ -38,38 +35,32 @@ def _fwd_kernel_token_att1(
 
     cur_kv_head = cur_head // kv_group_num
 
-    offs_d = tl.arange(0, BLOCK_DMODEL)  # [D]
+    offs_d = tl.arange(0, BLOCK_DMODEL)
     cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
-    cur_batch_in_all_start_index = tl.load(B_Att_Start_Loc + cur_batch)  # use window index
+    cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
     cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
-    cur_att_seq_len = tl.load(B_Att_Seqlen + cur_batch)
 
-    # use new start index of k value
-    cur_batch_start_index = tl.maximum(cur_batch_seq_len - sliding_window, 0)
+    cur_batch_start_index = 0
     cur_batch_end_index = cur_batch_seq_len
 
-    off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d * stride_qd  # [D]
+    off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d * stride_qd
 
-    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)  # [32]
+    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
 
-    # use new value to decide block mask
     block_stard_index = start_n * BLOCK_N
-    block_mask = tl.where(block_stard_index < cur_att_seq_len, 1, 0)  # a number
+    block_mask = tl.where(block_stard_index < cur_batch_seq_len, 1, 0)
 
     for start_mark in range(0, block_mask, 1):
-        q = tl.load(Q + off_q + start_mark)  # [SYM] why here add start_mark
-        offs_n_new = cur_batch_start_index + offs_n  # the latest window of token
+        q = tl.load(Q + off_q + start_mark)
+        offs_n_new = cur_batch_start_index + offs_n
         k_loc = tl.load(
             Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * offs_n_new,
             mask=offs_n_new < cur_batch_end_index,
             other=0,
-        )
-        off_k = (
-            k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :] * stride_kd
-        )  # [32, D], find token index
+        ).to(tl.int64)
+        off_k = k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :] * stride_kd
         k = tl.load(K + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0)
-        att_value = tl.sum(q[None, :] * k, 1)  # [1, D] * [32, D] = [32, D] -> [32]
-        att_value = att_value.to(tl.float32)
+        att_value = tl.sum(q[None, :] * k, 1, dtype=tl.float32)
         att_value *= sm_scale
         off_o = cur_head * att_stride_h + (cur_batch_in_all_start_index + offs_n) * att_stride_bs
         tl.store(Att_Out + off_o, att_value, mask=offs_n_new < cur_batch_end_index)
@@ -77,19 +68,17 @@ def _fwd_kernel_token_att1(
 
 
 @torch.no_grad()
-def token_att_fwd(
-    q, k, att_out, Req_to_tokens, B_req_idx, B_Start_Loc, B_Seqlen, B_Att_Start_Loc, B_Att_Seqlen, sliding_window
-):
+def token_att_fwd(q, k, att_out, Req_to_tokens, B_req_idx, B_Start_Loc, B_Seqlen, max_len_in_batch):
     BLOCK = 32
     # shape constraints
     Lq, Lk = q.shape[-1], k.shape[-1]
     assert Lq == Lk
-    assert Lk in {16, 32, 64, 128}
+    assert Lk in {16, 32, 64, 128, 256}
     sm_scale = 1.0 / (Lk ** 0.5)
 
     batch, head_num = B_req_idx.shape[0], q.shape[1]
 
-    grid = (batch, head_num, triton.cdiv(sliding_window, BLOCK))
+    grid = (batch, head_num, triton.cdiv(max_len_in_batch, BLOCK))
     kv_group_num = q.shape[1] // k.shape[1]
 
     if kv_group_num == 1:
@@ -105,8 +94,6 @@ def token_att_fwd(
         B_req_idx,
         B_Start_Loc,
         B_Seqlen,
-        B_Att_Start_Loc,
-        B_Att_Seqlen,
         att_out,
         Req_to_tokens.stride(0),
         Req_to_tokens.stride(1),
@@ -119,7 +106,6 @@ def token_att_fwd(
         att_out.stride(0),
         att_out.stride(1),
         kv_group_num=kv_group_num,
-        sliding_window=sliding_window,
         BLOCK_DMODEL=Lk,
         BLOCK_N=BLOCK,
         num_warps=num_warps,
diff --git a/lightllm/models/mistral/triton_kernel/token_attention_nopad_reduceV.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py
similarity index 59%
rename from lightllm/models/mistral/triton_kernel/token_attention_nopad_reduceV.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py
index acf4923f82..96a5b26dd6 100644
--- a/lightllm/models/mistral/triton_kernel/token_attention_nopad_reduceV.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py
@@ -13,8 +13,6 @@ def _fwd_kernel_token_att2(
     B_req_idx,
     B_Start_Loc,
     B_Seqlen,
-    B_Att_Start_Loc,
-    B_Att_Seqlen,
     stride_req_to_tokens_b,
     stride_req_to_tokens_s,
     stride_ph,
@@ -26,7 +24,6 @@ def _fwd_kernel_token_att2(
     stride_oh,
     stride_od,
     kv_group_num,
-    sliding_window,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
 ):
@@ -35,36 +32,30 @@ def _fwd_kernel_token_att2(
 
     cur_kv_head = cur_head // kv_group_num
 
-    offs_n = tl.arange(0, BLOCK_N)  # [64]
-    offs_d = tl.arange(0, BLOCK_DMODEL)  # [D]
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
     cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
-    cur_batch_start_index = tl.maximum(cur_batch_seq_len - sliding_window, 0)  # new index
-    # cur_batch_end_index = cur_batch_seq_len
-    cur_batch_in_all_start_index = tl.load(B_Att_Start_Loc + cur_batch)  # new index
+    cur_batch_start_index = 0
+    cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
     cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
-    cur_att_seq_len = tl.load(B_Att_Seqlen + cur_batch)  # att length
 
-    v_loc_off = (
-        cur_batch_req_idx * stride_req_to_tokens_b + (cur_batch_start_index + offs_n) * stride_req_to_tokens_s
-    )  # the latest window of value [64]
-    p_offs = cur_head * stride_ph + (cur_batch_in_all_start_index + offs_n) * stride_pbs  # [64]
-    v_offs = cur_kv_head * stride_vh + offs_d[None, :] * stride_vd  # [1, D]
+    v_loc_off = cur_batch_req_idx * stride_req_to_tokens_b + (cur_batch_start_index + offs_n) * stride_req_to_tokens_s
+    p_offs = cur_head * stride_ph + (cur_batch_in_all_start_index + offs_n) * stride_pbs
+    v_offs = cur_kv_head * stride_vh + offs_d[None, :] * stride_vd
 
-    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)  # [D]
-    for start_n in range(0, cur_att_seq_len, BLOCK_N):
-        start_n = tl.multiple_of(start_n, BLOCK_N)  # check
-        p_value = tl.load(Prob + p_offs + start_n, mask=(start_n + offs_n) < cur_att_seq_len, other=0.0)  # [64]
+    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
+    for start_n in range(0, cur_batch_seq_len, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        p_value = tl.load(Prob + p_offs + start_n, mask=(start_n + offs_n) < cur_batch_seq_len, other=0.0)
         v_loc = tl.load(
             Req_to_tokens + v_loc_off + start_n * stride_req_to_tokens_s,
-            mask=(start_n + offs_n + cur_batch_start_index) < cur_batch_seq_len,
+            mask=(start_n + offs_n) < cur_batch_seq_len,
             other=0.0,
-        )  # [64]
+        ).to(tl.int64)
         v_value = tl.load(
-            V + v_offs + v_loc[:, None] * stride_vbs,
-            mask=(start_n + offs_n[:, None] + cur_batch_start_index) < cur_batch_seq_len,
-            other=0.0,
-        )  # [1, D] + [64, 1] = [64, D]
-        acc += tl.sum(p_value[:, None] * v_value, 0)  # [64, 1] * [64, D] = [64, D] -> [D]
+            V + v_offs + v_loc[:, None] * stride_vbs, mask=(start_n + offs_n[:, None]) < cur_batch_seq_len, other=0.0
+        )
+        acc += tl.sum(p_value[:, None] * v_value, 0)
 
     acc = acc.to(Out.dtype.element_ty)
     off_o = cur_batch * stride_obs + cur_head * stride_oh + offs_d * stride_od
@@ -74,9 +65,7 @@ def _fwd_kernel_token_att2(
 
 
 @torch.no_grad()
-def token_att_fwd2(
-    prob, v, out, Req_to_tokens, B_req_idx, B_Start_Loc, B_Seqlen, B_Att_Start_Loc, B_Att_Seqlen, sliding_window
-):
+def token_att_fwd2(prob, v, out, Req_to_tokens, B_req_idx, B_Start_Loc, B_Seqlen):
     BLOCK = 128
     # BLOCK = 64 # for triton 2.0.0dev
     batch, head = B_req_idx.shape[0], prob.shape[0]
@@ -94,8 +83,6 @@ def token_att_fwd2(
         B_req_idx,
         B_Start_Loc,
         B_Seqlen,
-        B_Att_Start_Loc,
-        B_Att_Seqlen,
         Req_to_tokens.stride(0),
         Req_to_tokens.stride(1),
         prob.stride(0),
@@ -107,7 +94,6 @@ def token_att_fwd2(
         out.stride(1),
         out.stride(2),
         kv_group_num=kv_group_num,
-        siliding_window=sliding_window,
         BLOCK_DMODEL=dim,
         BLOCK_N=BLOCK,
         num_warps=num_warps,
diff --git a/lightllm/models/llama/triton_kernel/token_attention_nopad_softmax.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_softmax.py
similarity index 71%
rename from lightllm/models/llama/triton_kernel/token_attention_nopad_softmax.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_softmax.py
index 5e6040ac55..0bb6410e13 100644
--- a/lightllm/models/llama/triton_kernel/token_attention_nopad_softmax.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_softmax.py
@@ -5,11 +5,15 @@
 
 @triton.jit
 def _fwd_kernel_token_softmax(
-    Logics, B_Start_Loc, B_Seqlen,
+    Logics,
+    B_Start_Loc,
+    B_Seqlen,
     Prob_Out,
-    stride_logic_h, stride_logic_bs,
-    stride_prob_h, stride_prob_bs,
-    BLOCK_SIZE: tl.constexpr
+    stride_logic_h,
+    stride_logic_bs,
+    stride_prob_h,
+    stride_prob_bs,
+    BLOCK_SIZE: tl.constexpr,
 ):
     cur_batch = tl.program_id(0)
     cur_head = tl.program_id(1)
@@ -18,18 +22,25 @@ def _fwd_kernel_token_softmax(
     cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
     cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
 
-    row = tl.load(Logics + cur_head * stride_logic_h + (cur_batch_in_all_start_index + col_offsets) * stride_logic_bs,
-                  mask=col_offsets < cur_batch_seq_len, other=-float('inf')).to(tl.float32)
+    row = tl.load(
+        Logics + cur_head * stride_logic_h + (cur_batch_in_all_start_index + col_offsets) * stride_logic_bs,
+        mask=col_offsets < cur_batch_seq_len,
+        other=-float("inf"),
+    ).to(tl.float32)
 
     row_minus_max = row - tl.max(row, axis=0)
     numerator = tl.exp(row_minus_max)
     denominator = tl.sum(numerator, axis=0)
     softmax_output = numerator / denominator
 
-    tl.store(Prob_Out + cur_head * stride_prob_h + (cur_batch_in_all_start_index + col_offsets)
-             * stride_prob_bs, softmax_output, mask=col_offsets < cur_batch_seq_len)
+    tl.store(
+        Prob_Out + cur_head * stride_prob_h + (cur_batch_in_all_start_index + col_offsets) * stride_prob_bs,
+        softmax_output,
+        mask=col_offsets < cur_batch_seq_len,
+    )
     return
 
+
 @torch.no_grad()
 def token_softmax_fwd(Logics, B_Start_Loc, B_Seqlen, Prob_Out, max_input_len):
     BLOCK_SIZE = triton.next_power_of_2(max_input_len)
@@ -42,20 +53,26 @@ def token_softmax_fwd(Logics, B_Start_Loc, B_Seqlen, Prob_Out, max_input_len):
         num_warps = 16
 
     _fwd_kernel_token_softmax[(batch, head_num)](
-        Logics, B_Start_Loc, B_Seqlen,
+        Logics,
+        B_Start_Loc,
+        B_Seqlen,
         Prob_Out,
-        Logics.stride(0), Logics.stride(1),
-        Prob_Out.stride(0), Prob_Out.stride(1),
+        Logics.stride(0),
+        Logics.stride(1),
+        Prob_Out.stride(0),
+        Prob_Out.stride(1),
         num_warps=num_warps,
         BLOCK_SIZE=BLOCK_SIZE,
     )
     return
 
+
 def test1():
 
     import torch
 
     B, N_CTX, H, D = 4, 1025, 12, 128
+    del D
 
     dtype = torch.float16
 
diff --git a/lightllm/models/llama/triton_kernel/token_attention_softmax_and_reducev.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_softmax_and_reducev.py
similarity index 100%
rename from lightllm/models/llama/triton_kernel/token_attention_softmax_and_reducev.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_softmax_and_reducev.py
diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/ppl_fp16/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/ppl_fp16/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/llama/triton_kernel/ppl_fp16_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/ppl_fp16/ppl_fp16_flash_decoding.py
similarity index 51%
rename from lightllm/models/llama/triton_kernel/ppl_fp16_flash_decoding.py
rename to lightllm/common/basemodel/triton_kernel/att/decode_att/ppl_fp16/ppl_fp16_flash_decoding.py
index 8fda084605..b0a9b6245c 100644
--- a/lightllm/models/llama/triton_kernel/ppl_fp16_flash_decoding.py
+++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/ppl_fp16/ppl_fp16_flash_decoding.py
@@ -1,27 +1,27 @@
 import torch
+from lightllm.utils.light_utils import HAS_LIGHTLLM_KERNEL, light_ops
 
 
-def token_decode_attention_flash_decoding(
-    q, infer_state, q_head_num, head_dim, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty
-):
+def token_decode_attention_flash_decoding(q, infer_state, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty):
     BLOCK_SEQ = 256
     batch_size = infer_state.batch_size
-    max_len_in_batch = infer_state.max_len_in_batch
+    q_head_num = q.shape[1]
+    head_dim = q.shape[2]
+    max_kv_seq_len = infer_state.max_kv_seq_len
     calcu_shape1 = (batch_size, q_head_num, head_dim)
 
-    from lightllm_ppl_fp16_flashdecoding_kernel import fp16_flashdecoding_stage1
-    from .flash_decoding_stage2 import flash_decode_stage2
+    from ..mha.flash_decoding.flash_decoding_stage2 import flash_decode_stage2
 
     o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out
 
     mid_o = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1, head_dim], dtype=torch.float16, device="cuda"
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1, head_dim], dtype=torch.float16, device="cuda"
     )
     mid_o_logexpsum = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1], dtype=torch.float16, device="cuda"
+        [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1], dtype=torch.float16, device="cuda"
     )
 
-    fp16_flashdecoding_stage1(
+    light_ops.fp16_flashdecoding_stage1(
         BLOCK_SEQ,
         mid_o,
         mid_o_logexpsum,
@@ -32,7 +32,7 @@ def token_decode_attention_flash_decoding(
         infer_state.req_manager.req_to_token_indexs,
         infer_state.b_req_idx,
         infer_state.b_seq_len,
-        infer_state.max_len_in_batch,
+        infer_state.max_kv_seq_len,
     )
 
     flash_decode_stage2(mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ)
diff --git a/lightllm/common/basemodel/triton_kernel/att/prefill_att/__init__.py b/lightllm/common/basemodel/triton_kernel/att/prefill_att/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py b/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py
similarity index 81%
rename from lightllm/models/llama/triton_kernel/context_flashattention_nopad.py
rename to lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py
index e36c51b394..5ba6d0beb6 100644
--- a/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py
+++ b/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py
@@ -336,25 +336,24 @@ def context_attention_fwd_no_prompt_cache(q, k, v, o, b_start_loc, b_seq_len, ma
 
 
 @triton.jit
-def _fwd_kernel_int8kv(
+def _fwd_kernel_contiguous_kv(
     Q,
     K,
     V,
     sm_scale,
     Out,
     B_Start_Loc,
+    B_kv_start_loc,
     B_Seqlen,
     b_prompt_cache_len,
     stride_qbs,
     stride_qh,
     stride_qd,
-    stride_kb,
-    stride_kh,
     stride_ks,
+    stride_kh,
     stride_kd,
-    stride_vb,
-    stride_vh,
     stride_vs,
+    stride_vh,
     stride_vd,
     stride_obs,
     stride_oh,
@@ -374,6 +373,7 @@ def _fwd_kernel_int8kv(
     prompt_cache_len = tl.load(b_prompt_cache_len + cur_batch)
     cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
     cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) - prompt_cache_len
+    kv_start_loc = tl.load(B_kv_start_loc + cur_batch)
 
     block_start_loc = BLOCK_M * start_m
 
@@ -393,6 +393,9 @@ def _fwd_kernel_int8kv(
     l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
     acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
 
+    stride_ks = tl.cast(stride_ks, tl.int64)
+    stride_vs = tl.cast(stride_vs, tl.int64)
+
     block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0)
     block_end_loc = tl.minimum(block_start_loc + BLOCK_M + prompt_cache_len, cur_batch_seq_len + prompt_cache_len)
     # causal mask
@@ -405,8 +408,7 @@ def _fwd_kernel_int8kv(
         #     other=0,
         # )
         off_k = (
-            cur_batch * stride_kb
-            + (start_n + offs_n[None, :]) * stride_ks
+            (kv_start_loc + start_n + offs_n[None, :]) * stride_ks
             + cur_kv_head * stride_kh
             + offs_d[:, None] * stride_kd
         )
@@ -432,8 +434,7 @@ def _fwd_kernel_int8kv(
         #     other=0.0,
         # )
         off_v = (
-            cur_batch * stride_vb
-            + (start_n + offs_n[:, None]) * stride_vs
+            (kv_start_loc + start_n + offs_n[:, None]) * stride_vs
             + cur_kv_head * stride_vh
             + offs_d[None, :] * stride_vd
         )
@@ -455,7 +456,9 @@ def _fwd_kernel_int8kv(
 
 
 @torch.no_grad()
-def context_attention_fwd_ppl_int8kv(q, k, v, o, b_start_loc, b_seq_len, max_input_len, b_prompt_cache_len):
+def context_attention_fwd_contiguous_kv(
+    q, k, v, o, b_start_loc, b_kv_start_loc, b_seq_len, max_q_input_len, b_prompt_cache_len
+):
     BLOCK_M = 128 if not is_tesla() else 64
     # shape constraints
     Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
@@ -468,34 +471,33 @@ def context_attention_fwd_ppl_int8kv(q, k, v, o, b_start_loc, b_seq_len, max_inp
     batch, head = b_seq_len.shape[0], q.shape[1]
     kv_group_num = q.shape[1] // k.shape[1]
 
-    grid = lambda meta: (triton.cdiv(max_input_len, meta["BLOCK_M"]), batch * head, 1)
+    grid = lambda meta: (triton.cdiv(max_q_input_len, meta["BLOCK_M"]), batch * head, 1)
     BLOCK_N = BLOCK_M
     num_warps = 4 if Lk <= 64 else 8
     num_stages = 1
 
-    _fwd_kernel_int8kv[grid](
-        q,
-        k,
-        v,
-        sm_scale,
-        o,
-        b_start_loc,
-        b_seq_len,
-        b_prompt_cache_len,
-        q.stride(0),
-        q.stride(1),
-        q.stride(2),
-        k.stride(0),
-        k.stride(1),
-        k.stride(2),
-        k.stride(3),
-        v.stride(0),
-        v.stride(1),
-        v.stride(2),
-        v.stride(3),
-        o.stride(0),
-        o.stride(1),
-        o.stride(2),
+    _fwd_kernel_contiguous_kv[grid](
+        Q=q,
+        K=k,
+        V=v,
+        sm_scale=sm_scale,
+        Out=o,
+        B_Start_Loc=b_start_loc,
+        B_kv_start_loc=b_kv_start_loc,
+        B_Seqlen=b_seq_len,
+        b_prompt_cache_len=b_prompt_cache_len,
+        stride_qbs=q.stride(0),
+        stride_qh=q.stride(1),
+        stride_qd=q.stride(2),
+        stride_ks=k.stride(0),
+        stride_kh=k.stride(1),
+        stride_kd=k.stride(2),
+        stride_vs=v.stride(0),
+        stride_vh=v.stride(1),
+        stride_vd=v.stride(2),
+        stride_obs=o.stride(0),
+        stride_oh=o.stride(1),
+        stride_od=o.stride(2),
         kv_group_num=kv_group_num,
         H=head,
         BLOCK_DMODEL=Lk,
@@ -596,86 +598,5 @@ def test():
     assert torch.allclose(torch_o, o, atol=1e-2, rtol=0)
 
 
-def torch_context_attention_fwd2(q, k, v, o, b_start_loc, b_seq_len, b_prompt_cache_len):
-
-    batch = b_start_loc.shape[0]
-    k = k.transpose(1, 2)
-    v = v.transpose(1, 2)
-    for i in range(batch):
-        start_loc = b_start_loc[i]
-        seq_len = b_seq_len[i]
-        prompt_cache_len = b_prompt_cache_len[i]
-        cur_q = q[start_loc : start_loc + seq_len - prompt_cache_len, :, :]
-        cur_q = cur_q.clone().to(torch.float32)
-        cur_k = k[i, :seq_len, :]
-        cur_k = cur_k.clone().to(torch.float32)
-
-        cur_v = v[i, :seq_len, :]
-        cur_v = cur_v.clone().to(torch.float32)
-
-        cur_q = cur_q.transpose(0, 1)
-        cur_k = cur_k.transpose(0, 1)
-        cur_v = cur_v.transpose(0, 1)
-        dk = cur_q.shape[-1]
-
-        p = torch.matmul(cur_q, cur_k.transpose(-2, -1)) / torch.sqrt(torch.tensor(dk, dtype=torch.float32))
-
-        q_index = torch.arange(cur_q.shape[1]).unsqueeze(-1).to(p.device)
-        k_index = torch.arange(cur_k.shape[1]).unsqueeze(0).to(p.device)
-        mask = (q_index + prompt_cache_len >= k_index).int()
-        mask = mask.unsqueeze(0).expand(cur_q.shape[0], -1, -1)
-
-        p = p.masked_fill(mask == 0, float("-inf"))
-
-        s = F.softmax(p, dim=-1)
-
-        o[start_loc : start_loc + seq_len - prompt_cache_len, :, :] = torch.matmul(s, cur_v).transpose(0, 1)
-
-
-def test2():
-    import torch
-    import numpy as np
-
-    Z, H, N_CTX, D_HEAD = 16, 16, 2048, 128
-    dtype = torch.float16
-    prompt_cache_len = 0
-    q = torch.empty((Z * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.1, std=0.2)
-    kv = torch.empty((Z, 2 * H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.4, std=0.2)
-    k = kv[:, :H]
-    v = kv[:, H:]
-    # v = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.3, std=0.2)
-    o = torch.empty((Z * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.3, std=0.2)
-    torch_o = torch.empty((Z * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda").normal_(
-        mean=0.3, std=0.2
-    )
-    max_input_len = N_CTX
-    b_start_loc = torch.zeros((Z,), dtype=torch.int32, device="cuda")
-    b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda")
-    b_prompt_cache_len = torch.zeros(Z, dtype=torch.int32, device="cuda")
-
-    for i in range(Z):
-        b_seq_len[i] = N_CTX
-        if i != 0:
-            b_start_loc[i] = b_start_loc[i - 1] + N_CTX - prompt_cache_len
-        b_prompt_cache_len[i] = prompt_cache_len
-    torch_context_attention_fwd2(q, k, v, torch_o, b_start_loc, b_seq_len, b_prompt_cache_len)
-
-    import time
-
-    torch.cuda.synchronize()
-    a = time.time()
-    for i in range(1000):
-        context_attention_fwd_ppl_int8kv(q, k, v, o, b_start_loc, b_seq_len, max_input_len, b_prompt_cache_len)
-    torch.cuda.synchronize()
-    b = time.time()
-    # print(o.shape, torch_out.shape)
-    print((b - a))
-
-    print("max ", torch.max(torch.abs(torch_o - o)))
-    print("mean ", torch.mean(torch.abs(torch_o - o)))
-    assert torch.allclose(torch_o, o, atol=1e-2, rtol=0)
-
-
 if __name__ == "__main__":
     test()
-    test2()
diff --git a/lightllm/common/basemodel/triton_kernel/dequantize_gemm_int4.py b/lightllm/common/basemodel/triton_kernel/dequantize_gemm_int4.py
deleted file mode 100644
index 143d93b230..0000000000
--- a/lightllm/common/basemodel/triton_kernel/dequantize_gemm_int4.py
+++ /dev/null
@@ -1,649 +0,0 @@
-import time
-
-import torch
-
-import triton
-import triton.language as tl
-
-
-@triton.autotune(
-	configs=[
-        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=8),
-        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=8),
-        triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
-        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8), 
-    ],
-	key=['M', 'N', 'K', 'NO_GROUPS'],
-)
-@triton.jit
-def matmul4_kernel(
-	a_ptr, b_ptr, c_ptr,
-	scales_ptr, zeros_ptr,
-	M, N, K,
-	stride_am, stride_ak,
-	stride_bk, stride_bn,
-	stride_cm, stride_cn,
-	stride_scales_g, stride_scales_n,
-	stride_zeros_g, stride_zeros_n,
-	groupsize, NO_GROUPS: tl.constexpr,
-	BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
-	GROUP_SIZE_M: tl.constexpr,
-):
-    """
-    Compute the matrix multiplication C = A x B.
-    A is of shape (M, K) float16
-    B is of shape (K//8, N) int32
-    C is of shape (M, N) float16
-    scales is of shape (G, N) float16
-    zeros is of shape (G, N//8) int32
-    groupsize is an int specifying the size of groups for scales and zeros.
-    G is K // groupsize.
-    Set NO_GROUPS to groupsize == K, in which case G = 1 and the kernel is more efficient.
-    WARNING: This kernel assumes that K is a multiple of BLOCK_SIZE_K.
-    WARNING: This kernel assumes that N is a multiple of BLOCK_SIZE_N.
-    WARNING: This kernel assumes that groupsize is a multiple of BLOCK_SIZE_K.
-    """
-    bits = 4
-    infearure_per_bits = 8
-    pid = tl.program_id(axis=0)
-    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-    num_pid_k = tl.cdiv(K, BLOCK_SIZE_K)
-    num_pid_in_group = GROUP_SIZE_M * num_pid_n
-    group_id = pid // num_pid_in_group
-    first_pid_m = group_id * GROUP_SIZE_M
-    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
-    pid_m = first_pid_m + (pid % group_size_m)
-    pid_n = (pid % num_pid_in_group) // group_size_m    
-    offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)   # (BLOCK_SIZE_M, BLOCK_SIZE_K)
-    a_mask = (offs_am[:, None] < M)
-    # b_ptrs is set up such that it repeats elements along the K axis 8 times
-    b_ptrs = b_ptr + ((offs_k[:, None] // infearure_per_bits) * stride_bk + offs_bn[None, :] * stride_bn)   # (BLOCK_SIZE_K, BLOCK_SIZE_N)
-    scales_ptrs = scales_ptr + offs_bn * stride_scales_n   # (BLOCK_SIZE_N,)
-    # zeros_ptrs is set up such that it repeats elements along the N axis 8 times
-    zeros_ptrs = zeros_ptr + ((offs_bn // infearure_per_bits) * stride_zeros_n)   # (BLOCK_SIZE_N,)
-    # shifter is used to extract the 4 bits of each element in the 32-bit word from B and zeros
-    shifter = (offs_k % infearure_per_bits) * bits
-    zeros_shifter = (offs_bn % infearure_per_bits) * bits
-    # If G == 1, scales and zeros are the same for all K, so we can load them once
-    if NO_GROUPS:
-        # Fetch scales and zeros; these are per-outfeature and thus reused in the inner loop
-        scales = tl.load(scales_ptrs)  # (BLOCK_SIZE_N,)
-        zeros = tl.load(zeros_ptrs)  # (BLOCK_SIZE_N,), each element is repeated 8 times, int32	
-        # Unpack zeros
-        zeros = (zeros >> zeros_shifter) & 0xF  # (BLOCK_SIZE_N,) int32
-        # zeros = (zeros + 1) * scales  # (BLOCK_SIZE_N,) float16
-        zeros = zeros * scales
-    # Now calculate a block of output of shape (BLOCK_SIZE_M, BLOCK_SIZE_N)
-    # M is along the batch dimension, N is along the outfeatures dimension, K is along the infeatures dimension
-    # So this loop is along the infeatures dimension (K)
-    # It's calculating BLOCK_SIZE_M batches in parallel, and for each batch, BLOCK_SIZE_N outfeatures in parallel
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, num_pid_k):
-        a = tl.load(a_ptrs, mask=a_mask, other=0.)   # (BLOCK_SIZE_M, BLOCK_SIZE_K)
-        b = tl.load(b_ptrs)  # (BLOCK_SIZE_K, BLOCK_SIZE_N), but repeated
-        if not NO_GROUPS:
-            g_id = k // (groupsize // BLOCK_SIZE_K)
-            ptr = scales_ptrs + g_id * stride_scales_g
-            scales = tl.load(ptr)  # (BLOCK_SIZE_N,)
-            ptr = zeros_ptrs + g_id * stride_zeros_g   # (BLOCK_SIZE_N,)
-            zeros = tl.load(ptr)  # (BLOCK_SIZE_N,), each element is repeated 8 times, int32	
-            # Unpack zeros
-            zeros = (zeros >> zeros_shifter) & 0xF  # (BLOCK_SIZE_N,) int32
-            zeros = (zeros) * scales  # (BLOCK_SIZE_N,) float16	
-        # Now we need to unpack b (which is 4-bit values) into 32-bit values
-        b = (b >> shifter[:, None]) & 0xF  # Extract the 4-bit values
-        b = b * scales[None, :] - zeros[None, :]  # Scale and shift
-        # print("data type", a, b)
-        accumulator += tl.dot(a, b.to(a.dtype))
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += (BLOCK_SIZE_K // infearure_per_bits) * stride_bk  
-    c = accumulator.to(c_ptr.dtype.element_ty)  
-    # Store the result
-    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-    tl.store(c_ptrs, accumulator, mask=c_mask)
-
-
-def matmul_dequantize_int4_gptq(x: torch.FloatTensor, qweight: torch.IntTensor, scales: torch.FloatTensor, qzeros: torch.IntTensor, group_size, output=None) -> torch.FloatTensor:
-	"""
-	Compute the matrix multiplication C = A x B + bias.
-	Where B is quantized using GPTQ and groupsize = -1 into 4-bit values.
-
-	A is of shape (..., K) float16
-	qweight is of shape (K//8, N) int32
-	scales is of shape (G, N) float16
-	qzeros is of shape (G, N//8) int32
-	bias is of shape (1, N) float16
-
-	groupsize is the number of infeatures in each group.
-	G = K // groupsize
-
-	Returns C of shape (..., N) float16
-	"""
-	assert x.shape[-1] == (qweight.shape[0] * 8), "A must be a multiple of 8 in the last dimension"
-	assert x.is_contiguous(), "A must be contiguous"
-
-	M, K = x.shape
-	N = qweight.shape[1]
-	# This is based on the possible BLOCK_SIZE_Ks
-	# assert K % 16 == 0 and K % 32 == 0 and K % 64 == 0 and K % 128 == 0, "K must be a multiple of 16, 32, 64, and 128"
-	# # This is based on the possible BLOCK_SIZE_Ns
-	# assert N % 16 == 0 and N % 32 == 0 and N % 64 == 0 and N % 128 == 0 and N % 256 == 0, "N must be a multiple of 16, 32, 64, 128, and 256"
-	# # This is based on the possible BLOCK_SIZE_Ks
-	# assert groupsize % 32 == 0 and groupsize % 64 == 0 and groupsize % 128 == 0, "groupsize must be a multiple of 32, 64, and 128"
-
-	# output = torch.empty((M, N), device='cuda', dtype=torch.float16)
-	if output is None:
-		inplace = False
-		output = torch.empty((M, N), device=x.device, dtype=x.dtype)
-	else:
-		inplace = True
-
-	grid = lambda META: (
-		triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),
-	)
-	matmul4_kernel[grid](
-		x, qweight, output,
-		scales, qzeros,
-		M, N, K,
-		x.stride(0), x.stride(1),
-		qweight.stride(0), qweight.stride(1),
-		output.stride(0), output.stride(1),
-		scales.stride(0), scales.stride(1),
-		qzeros.stride(0), qzeros.stride(1),
-		group_size, group_size == K,
-    )
-	# return output
-	if not inplace:
-		return output
-
-
-@triton.autotune(
-	configs=[
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=8),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 512, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=8),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 512, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-	    
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=8),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
-		triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 512, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=8),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 512, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-		
- ],
-	key=['M', 'N', 'K'],
-    reset_to_zero=['c_ptr']
-)
-@triton.jit
-def matmul_kernel(
-    a_ptr, b_ptr, c_ptr,
-    bs_ptr, bzp_ptr,
-    M, N, K,
-    stride_am, stride_ak,
-    stride_bk, stride_bn,
-    stride_cm, stride_cn,
-    stride_bsk, stride_bsn,
-    stride_bzpk, stride_bzpn,
-    group_size,
-    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
-    GROUP_SIZE_M: tl.constexpr, SPLIT_K: tl.constexpr
-    ):
-    """
-    assert K % (BLOCK_SIZE_K * SPLIT_K) == 0
-    """
-    pid = tl.program_id(axis=0)
-    pid_sp_k = tl.program_id(axis=1)
-    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-    num_pid_k = tl.cdiv(K, BLOCK_SIZE_K)
-    num_pid_in_group = GROUP_SIZE_M * num_pid_n
-    group_id = pid // num_pid_in_group
-    first_pid_m = group_id * GROUP_SIZE_M
-    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
-    pid_m = first_pid_m + (pid % group_size_m)
-    pid_n = (pid % num_pid_in_group) // group_size_m    
-    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
-    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
-    offs_k = pid_sp_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
-
-    # [BLOCK_M, BLOCK_K]
-    a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
-    # [BLOCK_K, BLOCK_N] but repeated 8 times in N
-    b_ptrs = b_ptr + (offs_k[:, None] // 8) * stride_bk + offs_bn[None, :] * stride_bn
-    # tl.static_print("shape", a_ptrs, b_ptrs, bs_ptrs, bzp_ptrs)
-    # -----------------------------------------------------------
-    # Iterate to compute a block of the C matrix.
-    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
-    # of fp32 values for higher accuracy.
-    # `accumulator` will be converted back to fp16 after the loop.
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)):
-        # Load the next block of A and B.
-        # [BLOCK_K, BLOCK_N] but repeated group_size times in K 
-        bs_ptrs = bs_ptr + ((offs_k[:, None] + k * BLOCK_SIZE_K * SPLIT_K) // group_size) * stride_bsk \
-            + offs_bn[None, :] * stride_bsn
-        # [BLOCK_K, BLOCK_N] but repeated in K and N
-        bzp_ptrs = bzp_ptr + ((offs_k[:, None] + k * BLOCK_SIZE_K * SPLIT_K) // group_size) * stride_bzpk \
-            + (offs_bn[None, :] // 8) * stride_bzpn
-        b_shift_bits = (offs_k[:, None] % 8) * 4 # assert BLOCK_SIZE_K % 8 == 0
-        bzp_shift_bits = (offs_bn[None, :] % 8) * 4
-        a = tl.load(a_ptrs)
-        b = tl.load(b_ptrs)
-        bs = tl.load(bs_ptrs)
-        bzp = tl.load(bzp_ptrs)
-        # We accumulate along the K dimension.
-        int_b = (b >> b_shift_bits) & 0xF
-        int_bzp = (bzp >> bzp_shift_bits) & 0xF
-        b = ((int_b - int_bzp) * bs).to(a.dtype)
-        accumulator += tl.dot(a, b.to(a.dtype))
-        # Advance the ptrs to the next K block.
-        a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
-        b_ptrs += (BLOCK_SIZE_K * SPLIT_K * stride_bk // 8)  # assert BLOCK_SIZE_K % 8 == 0
-    # You can fuse arbitrary activation functions here
-    # while the accumulator is still in FP32!
-    c = accumulator.to(c_ptr.dtype.element_ty)
-    # -----------------------------------------------------------
-    # Write back the block of the output matrix C with masks.
-    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-    if SPLIT_K == 1:
-        tl.store(c_ptrs, c, mask=c_mask)
-    else:
-        tl.atomic_add(c_ptrs, c, mask=c_mask)
-
-
-def matmul_dequantize_int4_s2(x: torch.FloatTensor, qweight: torch.IntTensor, scales: torch.FloatTensor, qzeros: torch.IntTensor, group_size: int = 128, output=None) -> torch.FloatTensor:
-    """
-    """
-    assert x.is_contiguous(), "A must be contiguous"
-    assert qweight.is_contiguous(), "B must be contiguous"  
-    M, K = x.shape
-    N = scales.shape[1]
-    if output is None:
-        output = torch.zeros((M, N), device=x.device, dtype=x.dtype)  
-    grid = lambda META: (
-        triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),
-        META['SPLIT_K'],
-    )
-    matmul_kernel[grid](
-        x, qweight, output,
-        scales, qzeros,
-        M, N, K,
-        x.stride(0), x.stride(1),
-        qweight.stride(0), qweight.stride(1),
-        output.stride(0), output.stride(1),
-        scales.stride(0), scales.stride(1),
-        qzeros.stride(0), qzeros.stride(1),
-        group_size,
-    )
-    return output
-
-
-@triton.autotune(
-    configs=[
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32}, num_stages=5, num_warps=2),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64}, num_stages=5, num_warps=2),
-        triton.Config({'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32}, num_stages=5, num_warps=2),
-        triton.Config({'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32}, num_stages=5, num_warps=2),
-    ],
-    key=['K', 'N'],
-)
-@triton.jit
-def dequantize_kernel(
-    # Pointers to matrices
-    b_ptr, b_scale_ptr, b_zp_ptr, fpb_ptr,
-    # Matrix dimensions
-    K, N, group_size,
-    stride_bk, stride_bn,
-    stride_bsk, stride_bsn,
-    stride_bzpk, stride_bzpn,
-    stride_fpbk, stride_fpbn,
-    # Meta-parameters
-    BLOCK_SIZE_K: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
-):
-    """Dequantize tile [BLOCK_SIZE_K, BLOCK_SIZE_N] in full precision.
-    We should assert BLOCK_SIZE_N % 8 == 0.
-    weight[K // 8, N], scale[K // group_size, N], zp[K // group_size, N // group_size]
-    """
-    k_block_idx = tl.program_id(axis=0)
-    n_block_idx = tl.program_id(axis=1)
-    offs_k = k_block_idx * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
-    offs_n = n_block_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    fpb_offs = offs_k[:, None] * stride_fpbk + offs_n[None, :] * stride_fpbn
-    b_offs = (offs_k[:, None] // 8) * stride_bk + offs_n[None, :] * stride_bn
-    bzp_offs = (offs_k[:, None] // group_size) * stride_bzpk + (offs_n[None, :] // 8) * stride_bzpn
-    bs_offs = (offs_k[:, None] // group_size) * stride_bsk + offs_n[None, :] * stride_bsn
-    n_mask = offs_n[None, :] < N
-    k_mask = offs_k[:, None] < K
-    mask = n_mask & k_mask
-    int32_b = tl.load(b_ptr + b_offs, mask=mask, other=0.0)
-    zp_b = tl.load(b_zp_ptr + bzp_offs, mask=mask, other=0.0)
-    scale_b = tl.load(b_scale_ptr + bs_offs, mask=mask, other=0.0)
-    b_shift = (offs_k[:, None] % 8) * 4
-    bzp_shift = (offs_n[None, :] % 8) * 4
-    fp_weight = (((int32_b >> b_shift) & 0xF) - ((zp_b >> bzp_shift) & 0xF)) * scale_b
-    tl.store(fpb_ptr + fpb_offs, fp_weight, mask=mask)
-
-
-def dequantize_int4(b, b_scale, b_zero_point, device, dtype, group_size):
-    Kw, N = b.shape
-    K = Kw * 8
-    fp_b = torch.ones((K, N), device=device, dtype=dtype)
-    grid = lambda META: (
-        triton.cdiv(K, META['BLOCK_SIZE_K']),
-        triton.cdiv(N, META['BLOCK_SIZE_N']), 
-    )
-    dequantize_kernel[grid](
-        b, b_scale, b_zero_point, fp_b,
-        K, N, group_size,
-        b.stride(0), b.stride(1),
-        b_scale.stride(0), b_scale.stride(1),
-        b_zero_point.stride(0), b_zero_point.stride(1),
-        fp_b.stride(0), fp_b.stride(1)
-    )
-    return fp_b
-
-
-def matmul_dequantize_int4_s1(a, b, b_scale, b_zero_point, group_size=128, out=None):
-    """
-    Matmul dequantize int4 s1 dequantize weight to `fp_b` and do fp16 torch.mm,
-    this is for `prefill` stage, since weight size is fixed so is dequantize overhead,
-    perfill stage have more tokens to amortize dequant cost.
-    """
-    assert a.is_contiguous(), "Matrix A must be contiguous"
-    # assert b.is_contiguous(), "Matrix B must be contiguous"
-    M, K = a.shape
-    Kw, N = b.shape
-    if out is None:
-        # Allocates output.
-        out = torch.empty((M, N), device=a.device, dtype=a.dtype)
-    fp_b = dequantize_int4(b, b_scale, b_zero_point, a.device, a.dtype, group_size)
-    torch.mm(a, fp_b, out=out)
-    fp_b = None
-    return out
-
-
-def quantize_int4(weight, group_size=128, tp_rank=0):
-    # Weight shape: [H1 // 8, H2]
-    # Scale shape: [H1 // group_size, H2]
-    # zero_pint shape: [H1 // group_size, H2 // 8]
-
-    weight = weight.transpose(1, 0)
-    h1, h2 = weight.shape
-    assert h1 % 8 == 0 and h2 % 8 == 0, "H1 {} H2 {}".format(h1, h2)
-    assert h2 % group_size == 0, "H1 {} H2 {}".format(h1, h2)
-    weight = weight.contiguous().view(-1, group_size).cuda(tp_rank)
-    weight_max = weight.amax(-1, keepdim=True)
-    weight_max = torch.where(weight_max < 0, 0, weight_max)
-    weight_min = weight.amin(-1, keepdim=True)
-    weight_min = torch.where(weight_min > 0, 0, weight_min)
-    weight_range = weight_max - weight_min 
-    scale = weight_range / (2 ** 4 - 1)
-    zero_point = (-weight_min / scale).round().clamp(0, 15).to(torch.int32)
-    weight = (weight / scale + zero_point).round().clamp(0, 15).to(torch.int32).view(h1, h2)
-    int_weight = torch.empty(h1, h2 // 8).to(torch.int32).to(weight.device)
-    int_zero_point = torch.zeros(h1 // 8, h2 // group_size).to(torch.int32).to(weight.device)
-    zero_point = zero_point.view(h1, -1)
-    scale = scale.view(h1, -1)
-    # pack 8 int4 in an int32 number.
-    # Weight pack in row.
-    for pack in range(0, h2, 8):
-        for i in range(8):
-            int_weight[:, pack // 8] += weight[:, pack + i] << (i * 4)
-    # zero point pack in col.
-    for pack in range(0, h1, 8):
-        for i in range(8):
-            int_zero_point[pack // 8, :] += zero_point[pack + i, :] << (i * 4)
-    '''
-    fp_weight = torch.zeros(h1, h2).half().to(weight.device)
-    for pack in range(0, h1 // 8):
-        for i in range(8):
-            fp_weight[pack * 8 + i, :] = \
-                ((int_weight[pack, :] << (28 - i * 4) >> 28) + 16) % 16
-    print((fp_weight - weight).abs().sum())
-
-    fp_zp = torch.zeros(zero_point.shape).half().to(zero_point.device)
-    for pack in range(0, h1 // 8):
-        for i in range(8):
-            fp_zp[pack * 8 + i, :] = \
-                (int_zero_point[pack, :] >> (i * 4)) & 15
-
-    print((fp_zp - zero_point).abs().sum())
-    '''
-    weight = None
-    return int_weight.transpose(1, 0).contiguous(), scale.transpose(1, 0).contiguous(), int_zero_point.transpose(1, 0).contiguous(), group_size
-
-
-def unpack_int4(weight, scale, zp):
-    """
-    Test function to verify quantize int4 is correct.
-    Will not be used in model inference.
-    """
-    weight = weight.transpose(1, 0)
-    scale = scale.transpose(1, 0)
-    zp = zp.transpose(1, 0)
-    h1, h2 = weight.shape
-    group_size = h2 * 8 // scale.shape[1]
-    group_num = scale.shape[1]
-    fp_weight = torch.zeros(h1, h2 * 8).half().to(weight.device)
-    fp_zero_point = torch.zeros(h1, group_num).to(weight.device)
-    for pack in range(0, h2):
-        for i in range(8):
-            fp_weight[:, pack * 8 + i] = (weight[:, pack] >> (i * 4)) & 0xF
-    for pack in range(0, h1 // 8):
-        for i in range(8):
-            fp_zero_point[pack * 8 + i, :] = (zp[pack, :] >> (i * 4)) & 0xF
-    for g in range(group_num):
-        fp_weight[:, g * group_size:(g + 1) * group_size] = (fp_weight[:, g * group_size:(g + 1) * group_size] - \
-                                                             fp_zero_point[:, g].unsqueeze(1)) * scale[:, g].unsqueeze(1)
-    return fp_weight.transpose(1, 0)
-
-
-def test_int4(M, K, N):
-    import time
-
-    print("M: {} K: {} N: {}".format(M, K, N))
-    a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-    b = torch.randn((K, N), device='cuda', dtype=torch.float16)
-    int_b, b_scale, b_zero_point, _ = quantize_int4(b)
-    for _ in range(10):
-        triton_output = matmul_dequantize_int4_s1(a, int_b, b_scale, b_zero_point)
-    torch.cuda.synchronize()
-    iters = 512
-    t1 = time.time()
-    for _ in range(iters):
-        triton_output = matmul_dequantize_int4_s1(a, int_b, b_scale, b_zero_point)
-    torch.cuda.synchronize()
-    t2 = time.time()
-    triton_time = t2 - t1
-    print("Triton time cost", (t2 - t1))
-    for _ in range(10):
-        torch_output = torch.matmul(a, b)
-    torch.cuda.synchronize()
-    iters = 512
-    t1 = time.time()
-    for _ in range(iters):
-        torch_output = torch.matmul(a, b)
-    torch.cuda.synchronize()
-    t2 = time.time()
-    torch_time = t2 - t1
-    print("Torch time cost", (t2 - t1))
-    return triton_time, torch_time
-
-
-def test_correct_int4_s1(M=32, K=4096, N=4096):
-    group_size = 128
-    a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-    b = torch.randn((K, N), device='cuda', dtype=torch.float16)
-    int_b, b_scale, b_zero_point, _ = quantize_int4(b, group_size=group_size)
-    cos = torch.nn.CosineSimilarity(0)
-    fp_weight = dequantize_int4(int_b, b_scale, b_zero_point, a.device, a.dtype, group_size)
-    print("Quantize cos", cos(fp_weight.flatten().to(torch.float32), b.flatten().to(torch.float32)))
-    triton_output = matmul_dequantize_int4_s1(a, int_b, b_scale, b_zero_point, group_size)
-    torch_output = torch.matmul(a, b)
-    print(f"triton_output={triton_output}")
-    print(f"torch_output={torch_output}")
-    print("Output cos", cos(triton_output.flatten().to(torch.float32), torch_output.flatten().to(torch.float32)))
-
-
-def test_correct_int4_s2(M=32, K=4096, N=4096):
-    group_size = 128
-    a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-    b = torch.randn((K, N), device='cuda', dtype=torch.float16)
-    int_b, b_scale, b_zero_point, _ = quantize_int4(b, group_size=group_size)
-    cos = torch.nn.CosineSimilarity(0)
-    fp_weight = unpack_int4(int_b, b_scale, b_zero_point)
-    print("Quantize cos", cos(fp_weight.flatten().to(torch.float32), b.flatten().to(torch.float32)))
-    triton_output = matmul_dequantize_int4_s2(a, int_b, b_scale, b_zero_point, group_size)
-    torch_output = torch.matmul(a, b)
-    print(f"triton_output={triton_output}")
-    print(f"torch_output={torch_output}")
-    print("Output cos", cos(triton_output.flatten().to(torch.float32), torch_output.flatten().to(torch.float32)))
-
-
-def test_correct_int4_gptq(M=32, K=4096, N=4096):
-    group_size = 128
-    a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-    b = torch.randn((K, N), device='cuda', dtype=torch.float16)
-    int_b, b_scale, b_zero_point, _ = quantize_int4(b, group_size=group_size)
-    cos = torch.nn.CosineSimilarity(0)
-    fp_weight = unpack_int4(int_b, b_scale, b_zero_point)
-    print("Quantize cos", cos(fp_weight.flatten().to(torch.float32), b.flatten().to(torch.float32)))
-    triton_output = matmul_dequantize_int4_gptq(a, int_b, b_scale, b_zero_point, group_size)
-    torch_output = torch.matmul(a, b)
-    print(f"triton_output={triton_output}")
-    print(f"torch_output={torch_output}")
-    print("Output cos", cos(triton_output.flatten().to(torch.float32), torch_output.flatten().to(torch.float32)))
-
-
-@triton.testing.perf_report(
-    triton.testing.Benchmark(
-        x_names=['M'],  # Argument names to use as an x-axis for the plot
-        x_vals=[4, 8, 16, 32, 64, 128] + [
-            128 * i for i in range(2, 33, 2)
-        ],  # Different possible values for `x_name`
-        line_arg='provider',  # Argument name whose value corresponds to a different line in the plot
-        # Possible values for `line_arg`
-        line_vals=['cublas', 'triton-s1', 'dequantize', 'triton-s2', 'triton-gptq'],
-        # Label name for the lines
-        line_names=["cuBLAS", "Triton-s1", "Dequant(GB/s)", "Triton-s2", "Triton-gptq"],
-        # Line styles
-        styles=[('green', '-'), ('blue', '-'), ('red', '-'), ('purple', '-'), ('yellow', '-')],
-        ylabel="TFLOPS",  # Label name for the y-axis
-        plot_name="matmul-performance",  # Name for the plot, used also as a file name for saving the plot.
-        args={},
-    )
-)
-def benchmark(M, provider):
-    K = 4096
-    N = 4096
-    a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-    b = torch.randn((K, N), device='cuda', dtype=torch.float16)
-    quantiles = [0.5, 0.2, 0.8]
-    if provider == 'cublas':
-        ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch.matmul(a, b), quantiles=quantiles)
-        perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
-    if provider == 'triton-s1':
-        intb, b_scale, bzp, _ = quantize_int4(b, group_size=64)
-        ms, min_ms, max_ms = triton.testing.do_bench(lambda: matmul_dequantize_int4_s1(a, intb, b_scale, bzp, 64), quantiles=quantiles)
-        perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
-    if provider == 'triton-s2':
-        intb, b_scale, bzp, _ = quantize_int4(b, group_size=64)
-        ms, min_ms, max_ms = triton.testing.do_bench(lambda: matmul_dequantize_int4_s2(a, intb, b_scale, bzp, 64), quantiles=quantiles)
-        perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
-    if provider == 'dequantize':
-        intb, b_scale, bzp, _ = quantize_int4(b, group_size=64)
-        ms, min_ms, max_ms = triton.testing.do_bench(lambda: dequantize_int4(intb, b_scale, bzp, 'cuda', torch.float16, 64), quantiles=quantiles)        
-        perf = lambda ms: 2 * M * K * 1e-9 / (ms * 1e-3)
-    if provider == 'triton-gptq':
-        intb, b_scale, bzp, _ = quantize_int4(b, group_size=64)
-        ms, min_ms, max_ms = triton.testing.do_bench(lambda: matmul_dequantize_int4_gptq(a, intb, b_scale, bzp, 64), quantiles=quantiles)
-        perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
-    return perf(ms), perf(max_ms), perf(min_ms)
-
-
-def test_model_layer(bs, sqe_len, hidden, inter, tp):
-    st1 = 0
-    st2 = 0
-    t1, t2 = test_int4(bs * sqe_len, hidden, hidden * 3 // tp)
-    st1 += t1
-    st2 += t2
-    t1, t2 = test_int4(bs * sqe_len, hidden // tp, hidden)
-    st1 += t1
-    st2 += t2
-    t1, t2 = test_int4(bs * sqe_len, hidden, inter * 2 // tp)
-    st1 += t1
-    st2 += t2
-    t1, t2 = test_int4(bs * sqe_len, inter // tp, hidden)
-    st1 += t1
-    st2 += t2
-    print("Triton time {} Torch time {}".format(st1, st2))
-
-
-if __name__ == "__main__":
-    # test_correct_int4_s1()
-    # test_correct_int4_s2()
-    # test_correct_int4_gptq()
-    benchmark.run(show_plots=True, print_data=True)
-    exit()
-    bs = 32
-    hidden = 4096
-    inter = 11008
-    prefill_len = 512
-    decode_len = 1
-    tp = 1
-    test_model_layer(bs, prefill_len, hidden, inter, tp)
-    test_model_layer(bs, decode_len, hidden, inter, tp)
diff --git a/lightllm/common/basemodel/triton_kernel/dequantize_gemm_int8.py b/lightllm/common/basemodel/triton_kernel/dequantize_gemm_int8.py
deleted file mode 100644
index e2c5c0dc95..0000000000
--- a/lightllm/common/basemodel/triton_kernel/dequantize_gemm_int8.py
+++ /dev/null
@@ -1,209 +0,0 @@
-import torch
-
-import triton
-import triton.language as tl
-
-
-@triton.autotune(
-    configs=[
-        triton.Config({'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128}, num_stages=3, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 256}, num_stages=3, num_warps=8),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32}, num_stages=5, num_warps=2),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64}, num_stages=5, num_warps=2),
-        triton.Config({'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64}, num_stages=3, num_warps=8),
-        triton.Config({'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32}, num_stages=4, num_warps=4),
-        triton.Config({'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32}, num_stages=5, num_warps=2),
-        triton.Config({'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32}, num_stages=5, num_warps=2),
-    ],
-    key=['K', 'N'],
-)
-
-
-@triton.jit
-def dequantize_kernel(
-    # Pointers to matrices
-    b_ptr, b_scale_ptr, fpb_ptr,
-    # Matrix dimensions
-    K, N,
-    stride_bk, stride_bn,
-    stride_fpbk, stride_fpbn,
-    # Meta-parameters
-    BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
-):
-    """Kernel for computing the matmul C = A x B.
-    A has shape (M, K), B has shape (K, N) and C has shape (M, N)
-    """
-    k_block_idx = tl.program_id(axis=0)
-    n_block_idx = tl.program_id(axis=1)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    offs_n = tl.arange(0, BLOCK_SIZE_N)
-    b_offs = (k_block_idx * BLOCK_SIZE_K + offs_k[:, None]) * stride_bk + \
-        (n_block_idx * BLOCK_SIZE_N + offs_n[None, :]) * stride_bn
-    fpb_offs = (k_block_idx * BLOCK_SIZE_K + offs_k[:, None]) * stride_fpbk + \
-        (n_block_idx * BLOCK_SIZE_N + offs_n[None, :]) * stride_fpbn
-    bs_offs = n_block_idx * BLOCK_SIZE_N + offs_n[None, :]
-    n_mask = n_block_idx * BLOCK_SIZE_N + offs_n[None, :] < N
-    mask = (k_block_idx * BLOCK_SIZE_K + offs_k[:, None] < K) & n_mask
-    int_b = tl.load(b_ptr + b_offs, mask=mask, other=0.0)
-    scale_b = tl.load(b_scale_ptr + bs_offs, mask=n_mask, other=0.0)
-    tl.store(fpb_ptr + fpb_offs, int_b * scale_b, mask=mask)
-
-
-def matmul_dequantize_int8(a, b, b_scale, out=None):
-    # Check constraints.
-    assert a.shape[1] == b.shape[0], "Incompatible dimensions"
-    assert a.is_contiguous(), "Matrix A must be contiguous"
-    # assert b.is_contiguous(), "Matrix B must be contiguous"
-    M, K = a.shape
-    K, N = b.shape
-    if out == None:
-        # Allocates output.
-        c = torch.empty((M, N), device=a.device, dtype=a.dtype)
-    else:
-        c = out
-    fp_b = torch.empty((K, N), device=a.device, dtype=a.dtype)
-    grid = lambda META: (
-        triton.cdiv(K, META['BLOCK_SIZE_K']), triton.cdiv(N, META['BLOCK_SIZE_N']),
-    )
-    dequantize_kernel[grid](
-        b, b_scale, fp_b,
-        K, N,
-        b.stride(0), b.stride(1),
-        fp_b.stride(0), fp_b.stride(1)
-    )
-    torch.mm(a, fp_b, out=c)
-    return c
-
-
-def quantize_int8(weight, axis=0, tp_rank=0):
-    # Weight shape: [H1, H2]
-    # Scale shape: [H2]
-    scale = weight.abs().amax(axis, keepdim=True) / 127.
-    weight = (weight / scale).to(torch.int8)
-    if axis == 0:
-        weight = weight.t().contiguous().t()
-    scale = scale.squeeze(axis)
-    return weight.contiguous().cuda(tp_rank), scale.contiguous().cuda(tp_rank)
-
-
-def test_int8(M, K, N):
-    import time
-
-    print("M: {} K: {} N: {}".format(M, K, N))
-    torch.manual_seed(0)
-    a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-    b = torch.randn((K, N), device='cuda', dtype=torch.float16)
-    int_b, b_scale = quantize_int8(b)
-    for _ in range(10):
-        triton_output = matmul_dequantize_int8(a, int_b, b_scale.unsqueeze(0))
-    torch.cuda.synchronize()
-    iters = 512
-    t1 = time.time()
-    for _ in range(iters):
-        triton_output = matmul_dequantize_int8(a, int_b, b_scale.unsqueeze(0))
-    torch.cuda.synchronize()
-    t2 = time.time()
-    triton_time = t2 - t1
-    print("Triton time cost", (t2 - t1))
-    for _ in range(10):
-        torch_output = torch.matmul(a, b)
-    torch.cuda.synchronize()
-    iters = 512
-    t1 = time.time()
-    for _ in range(iters):
-        torch_output = torch.matmul(a, b)
-    torch.cuda.synchronize()
-    t2 = time.time()
-    torch_time = t2 - t1
-    print("Torch time cost", (t2 - t1))
-    return triton_time, torch_time
-
-
-def test_correct_int8(M=512, K=4096, N=4096):
-    import time
-
-    a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-    b = torch.randn((K, N), device='cuda', dtype=torch.float16)
-    int_b, b_scale = quantize_int8(b)
-    cos = torch.nn.CosineSimilarity(0)
-    triton_output = matmul_dequantize_int8(a, int_b, b_scale)
-    torch_output = torch.matmul(a, b)
-    print(f"triton_output={triton_output}")        
-    print(f"torch_output={torch_output}")
-    print("Output cos ", cos(triton_output.flatten().to(torch.float32), torch_output.flatten().to(torch.float32)))
-
-
-@triton.testing.perf_report(
-    triton.testing.Benchmark(
-        x_names=['M', 'N', 'K'],  # Argument names to use as an x-axis for the plot
-        x_vals=[32, 64, 128, 256] + [
-            512 * i for i in range(1, 33)
-        ],  # Different possible values for `x_name`
-        line_arg='provider',  # Argument name whose value corresponds to a different line in the plot
-        # Possible values for `line_arg`
-        line_vals=['cublas', 'triton'],
-        # Label name for the lines
-        line_names=["cuBLAS", "Triton"],
-        # Line styles
-        styles=[('green', '-'), ('blue', '-')],
-        ylabel="TFLOPS",  # Label name for the y-axis
-        plot_name="matmul-performance",  # Name for the plot, used also as a file name for saving the plot.
-        args={},
-    )
-)
-
-
-def benchmark(M, N, K, provider):
-    quantiles = [0.5, 0.2, 0.8]
-    if provider == 'cublas':
-        a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-        b = torch.randn((K, N), device='cuda', dtype=torch.float16)
-        ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch.matmul(a, b), quantiles=quantiles)
-    if provider == 'triton':
-        a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-        b = torch.randn((K, N), device='cuda', dtype=torch.float16)
-        intb, b_scale = quantize_int8(b)
-        ms, min_ms, max_ms = triton.testing.do_bench(lambda: matmul_dequantize_int8(a, intb, b_scale), quantiles=quantiles)
-    perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
-    return perf(ms), perf(min_ms), perf(max_ms)
-
-
-def test_model_layer(bs, sqe_len, hidden, inter, tp):
-    st1 = 0
-    st2 = 0
-    t1, t2 = test_int8(bs * sqe_len, hidden, hidden * 3 // tp)
-    st1 += t1
-    st2 += t2
-    t1, t2 = test_int8(bs * sqe_len, hidden // tp, hidden)
-    st1 += t1
-    st2 += t2
-    t1, t2 = test_int8(bs * sqe_len, hidden, inter * 2 // tp)
-    st1 += t1
-    st2 += t2
-    t1, t2 = test_int8(bs * sqe_len, inter // tp, hidden)
-    st1 += t1
-    st2 += t2
-    print("Triton time {} Torch time {}".format(st1, st2))
-
-
-if __name__ == "__main__":
-    test_correct_int8()
-    benchmark.run(show_plots=True, print_data=True)
-
-    bs = 32
-    hidden = 4096
-    inter  = 11008
-    prefill_len = 512
-    decode_len = 1
-    tp = 1
-    test_model_layer(bs, prefill_len, hidden, inter, tp)
-    test_model_layer(bs, decode_len, hidden, inter, tp)
\ No newline at end of file
diff --git a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py
index bd53de386e..060a92bf7c 100644
--- a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py
+++ b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py
@@ -16,12 +16,13 @@ def _fwd_kernel_destindex_copy_kv(
     stride_o_h,
     stride_o_d,
     head_num,
+    head_dim,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_HEAD: tl.constexpr,
 ):
     cur_index = tl.program_id(0)
     offs_h = tl.arange(0, BLOCK_HEAD)
-    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_d = (tl.arange(0, BLOCK_DMODEL)) % head_dim
 
     dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
 
@@ -54,133 +55,10 @@ def destindex_copy_kv(K, DestLoc, Out):
         Out.stride(1),
         Out.stride(2),
         head_num,
-        BLOCK_DMODEL=head_dim,
+        head_dim,
+        BLOCK_DMODEL=triton.next_power_of_2(head_dim),
         BLOCK_HEAD=BLOCK_HEAD,
         num_warps=num_warps,
         num_stages=1,
     )
     return
-
-
-@triton.jit
-def _fwd_kernel_destindex_copy_quantize_kv(
-    K,
-    Dest_loc,
-    Out,
-    Out_scale,
-    stride_k_bs,
-    stride_k_h,
-    stride_k_d,
-    stride_o_bs,
-    stride_o_h,
-    stride_o_d,
-    stride_os_bs,
-    stride_os_h,
-    stride_os_d,
-    head_num,
-    BLOCK_DMODEL: tl.constexpr,
-    BLOCK_HEAD: tl.constexpr,
-):
-    cur_index = tl.program_id(0)
-    offs_h = tl.arange(0, BLOCK_HEAD)
-    offs_d = tl.arange(0, BLOCK_DMODEL)
-
-    dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
-    src_data = tl.load(
-        K + cur_index * stride_k_bs + offs_h[:, None] * stride_k_h + stride_k_d * offs_d[None, :],
-        mask=offs_h[:, None] < head_num,
-        other=0.0,
-    )
-    abs_data = tl.abs(src_data)
-    data_scale = (tl.max(abs_data, axis=1) / 127.0).to(Out_scale.dtype.element_ty)[:, None]
-    q_src_data = (src_data / data_scale).to(tl.int8)
-    o_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :]
-    os_ptrs = Out_scale + dest_index * stride_os_bs + stride_os_h * offs_h[:, None]
-    tl.store(o_ptrs, q_src_data, mask=offs_h[:, None] < head_num)
-    tl.store(os_ptrs, data_scale, mask=offs_h[:, None] < head_num)
-
-
-@torch.no_grad()
-def destindex_copy_quantize_kv(K, DestLoc, Out, Out_scale):
-    seq_len = DestLoc.shape[0]
-    head_num = K.shape[1]
-    head_dim = K.shape[2]
-    assert K.shape[1] == Out.shape[1] and K.shape[2] == Out.shape[2]
-    BLOCK_HEAD = triton.next_power_of_2(head_num)
-    grid = (seq_len,)
-    num_warps = 1
-
-    _fwd_kernel_destindex_copy_quantize_kv[grid](
-        K,
-        DestLoc,
-        Out,
-        Out_scale,
-        K.stride(0),
-        K.stride(1),
-        K.stride(2),
-        Out.stride(0),
-        Out.stride(1),
-        Out.stride(2),
-        Out_scale.stride(0),
-        Out_scale.stride(1),
-        Out_scale.stride(2),
-        head_num,
-        BLOCK_DMODEL=head_dim,
-        BLOCK_HEAD=BLOCK_HEAD,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
-
-
-def test1():
-    import time
-
-    B, N_CTX, H, D = 32, 1024, 12, 128
-    dest = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda()
-    src = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda()
-    dest_loc = torch.arange(0, B * N_CTX, dtype=torch.int32, device="cuda")
-
-    for _ in range(10):
-        destindex_copy_kv(src, dest_loc, dest)
-    torch.cuda.synchronize()
-    t1 = time.time()
-    for _ in range(1000):
-        destindex_copy_kv(src, dest_loc, dest)
-    torch.cuda.synchronize()
-    t2 = time.time()
-
-    print("Time cost ", t2 - t1)
-    print("max ", torch.max(torch.abs(dest - src)))
-    print("mean ", torch.mean(torch.abs(dest - src)))
-    assert torch.allclose(src, dest, atol=1e-2, rtol=0)
-
-
-def test2():
-    import time
-
-    B, N_CTX, H, D = 32, 1024, 12, 128
-    src = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda()
-    dest_loc = torch.arange(0, B * N_CTX, dtype=torch.int32).cuda()
-    value_dest = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda().to(torch.int8)
-    scale_dest = torch.randn((B * N_CTX, H, 1), dtype=torch.float16).cuda()
-
-    for _ in range(10):
-        destindex_copy_quantize_kv(src, dest_loc, value_dest, scale_dest)
-    torch.cuda.synchronize()
-    t1 = time.time()
-    for _ in range(1000):
-        destindex_copy_quantize_kv(src, dest_loc, value_dest, scale_dest)
-    torch.cuda.synchronize()
-    t2 = time.time()
-
-    print("Time cost ", t2 - t1)
-    print("max ", torch.max(torch.abs(value_dest * scale_dest - src)))
-    print("mean ", torch.mean(torch.abs(value_dest * scale_dest - src)))
-    cos = torch.nn.CosineSimilarity(0)
-    print("cos ", cos(src.flatten().to(torch.float32), (value_dest * scale_dest).flatten().to(torch.float32)))
-
-
-if __name__ == "__main__":
-    test1()
-    test2()
diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/__init__.py b/lightllm/common/basemodel/triton_kernel/fused_moe/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/common/fused_moe/deepep_scatter_gather.py b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_scatter_gather.py
similarity index 100%
rename from lightllm/common/fused_moe/deepep_scatter_gather.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/deepep_scatter_gather.py
diff --git a/lightllm/common/fused_moe/grouped_fused_moe.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py
similarity index 99%
rename from lightllm/common/fused_moe/grouped_fused_moe.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py
index 758d83ba34..f29d3a2a0e 100644
--- a/lightllm/common/fused_moe/grouped_fused_moe.py
+++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py
@@ -28,7 +28,7 @@
 from .moe_kernel_configs import MoeGroupedGemmKernelConfig
 from .moe_silu_and_mul import silu_and_mul_fwd
 from .moe_sum_reduce import moe_sum_reduce
-from lightllm.common.quantization.triton_quant.fp8.fp8act_quant_kernel import per_token_group_quant_fp8
+from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8
 from lightllm.utils.torch_ops_utils import direct_register_custom_op
 from lightllm.common.triton_utils.autotuner import autotune
 
diff --git a/lightllm/common/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py
similarity index 96%
rename from lightllm/common/fused_moe/grouped_fused_moe_ep.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py
index 5cc0d7a9be..2c6d013bd5 100644
--- a/lightllm/common/fused_moe/grouped_fused_moe_ep.py
+++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py
@@ -6,13 +6,15 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch.distributed as dist
 from lightllm.utils.log_utils import init_logger
-from lightllm.common.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
-from lightllm.common.fused_moe.moe_silu_and_mul_mix_quant_ep import silu_and_mul_masked_post_quant_fwd
-from lightllm.common.quantization.triton_quant.fp8.fp8act_quant_kernel import (
+from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
+from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul_mix_quant_ep import (
+    silu_and_mul_masked_post_quant_fwd,
+)
+from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import (
     per_token_group_quant_fp8,
     tma_align_input_scale,
 )
-from lightllm.common.fused_moe.deepep_scatter_gather import ep_scatter, ep_gather
+from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_scatter_gather import ep_scatter, ep_gather
 from lightllm.utils.envs_utils import get_deepep_num_max_dispatch_tokens_per_rank
 from lightllm.common.triton_utils.autotuner import Autotuner
 import numpy as np
diff --git a/lightllm/common/fused_moe/grouped_topk.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py
similarity index 100%
rename from lightllm/common/fused_moe/grouped_topk.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py
diff --git a/lightllm/common/fused_moe/moe_kernel_configs.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_kernel_configs.py
similarity index 100%
rename from lightllm/common/fused_moe/moe_kernel_configs.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/moe_kernel_configs.py
diff --git a/lightllm/common/fused_moe/moe_silu_and_mul.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py
similarity index 100%
rename from lightllm/common/fused_moe/moe_silu_and_mul.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py
diff --git a/lightllm/common/fused_moe/moe_silu_and_mul_config.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_config.py
similarity index 100%
rename from lightllm/common/fused_moe/moe_silu_and_mul_config.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_config.py
diff --git a/lightllm/common/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py
similarity index 100%
rename from lightllm/common/fused_moe/moe_silu_and_mul_mix_quant_ep.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py
diff --git a/lightllm/common/fused_moe/moe_sum_recude_config.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_sum_recude_config.py
similarity index 100%
rename from lightllm/common/fused_moe/moe_sum_recude_config.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/moe_sum_recude_config.py
diff --git a/lightllm/common/fused_moe/moe_sum_reduce.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_sum_reduce.py
similarity index 100%
rename from lightllm/common/fused_moe/moe_sum_reduce.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/moe_sum_reduce.py
diff --git a/lightllm/common/fused_moe/softmax_topk.py b/lightllm/common/basemodel/triton_kernel/fused_moe/softmax_topk.py
similarity index 100%
rename from lightllm/common/fused_moe/softmax_topk.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/softmax_topk.py
diff --git a/lightllm/common/fused_moe/topk_select.py b/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py
similarity index 96%
rename from lightllm/common/fused_moe/topk_select.py
rename to lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py
index 5206800efc..72c3a381ed 100644
--- a/lightllm/common/fused_moe/topk_select.py
+++ b/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py
@@ -22,7 +22,7 @@
 from lightllm.utils.sgl_utils import sgl_ops
 from lightllm.utils.light_utils import light_ops
 from typing import Callable, List, Optional, Tuple
-from lightllm.common.fused_moe.softmax_topk import softmax_topk
+from lightllm.common.basemodel.triton_kernel.fused_moe.softmax_topk import softmax_topk
 from lightllm.common.triton_utils.autotuner import Autotuner
 
 use_cuda_grouped_topk = os.getenv("LIGHTLLM_CUDA_GROUPED_TOPK", "False").upper() in ["ON", "TRUE", "1"]
@@ -177,8 +177,8 @@ def select_experts(
     scoring_func: str = "softmax",
     custom_routing_function: Optional[Callable] = None,
 ):
-    from lightllm.common.fused_moe.topk_select import fused_topk
-    from lightllm.common.fused_moe.grouped_topk import triton_grouped_topk
+    from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import fused_topk
+    from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import triton_grouped_topk
 
     # DeekSeekv2 uses grouped_top_k
     if use_grouped_topk:
diff --git a/lightllm/common/basemodel/triton_kernel/gen_decode_params.py b/lightllm/common/basemodel/triton_kernel/gen_decode_params.py
index 9804e46681..c8a6a850bc 100644
--- a/lightllm/common/basemodel/triton_kernel/gen_decode_params.py
+++ b/lightllm/common/basemodel/triton_kernel/gen_decode_params.py
@@ -9,8 +9,6 @@
 def gen_decode_params(b_seq_len: torch.Tensor):
     b_kv_seq_len = b_seq_len
     position_ids = b_seq_len - 1
-    mtp_step = get_env_start_args().mtp_step
-    mtp_size = mtp_step + 1
-    b_q_seq_len = torch.ones(b_seq_len.shape[0] // mtp_size, dtype=torch.int32, device=b_seq_len.device) * mtp_size
-    b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len[mtp_size - 1 :: mtp_size])
+    b_q_seq_len = torch.ones(b_seq_len.shape[0], dtype=torch.int32, device=b_seq_len.device)
+    b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len)
     return b_q_seq_len, b1_cu_q_seq_len, b_kv_seq_len, b1_cu_kv_seq_len, position_ids
diff --git a/lightllm/common/basemodel/triton_kernel/gen_prefill_params.py b/lightllm/common/basemodel/triton_kernel/gen_prefill_params.py
index e73b342994..8f9172b552 100644
--- a/lightllm/common/basemodel/triton_kernel/gen_prefill_params.py
+++ b/lightllm/common/basemodel/triton_kernel/gen_prefill_params.py
@@ -43,6 +43,7 @@ def _gen_cumsum_pad0_kernel(
 def gen_cumsum_pad0_tensor(b_q_seq_len: torch.Tensor, b_kv_seq_len: torch.Tensor):
     assert len(b_q_seq_len.shape) == 1
     assert b_q_seq_len.shape == b_kv_seq_len.shape
+    assert b_q_seq_len.is_contiguous()
 
     b1_cu_q_seq_len = torch.empty((b_q_seq_len.shape[0] + 1,), dtype=torch.int32, device="cuda")
     b1_cu_kv_seq_len = torch.empty((b_kv_seq_len.shape[0] + 1,), dtype=torch.int32, device="cuda")
diff --git a/lightllm/common/basemodel/triton_kernel/kv_copy/__init__.py b/lightllm/common/basemodel/triton_kernel/kv_copy/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/deepseek2/triton_kernel/destindex_copy_kv.py b/lightllm/common/basemodel/triton_kernel/kv_copy/mla_copy_kv.py
similarity index 92%
rename from lightllm/models/deepseek2/triton_kernel/destindex_copy_kv.py
rename to lightllm/common/basemodel/triton_kernel/kv_copy/mla_copy_kv.py
index 39deb1b6f7..41a25877a7 100644
--- a/lightllm/models/deepseek2/triton_kernel/destindex_copy_kv.py
+++ b/lightllm/common/basemodel/triton_kernel/kv_copy/mla_copy_kv.py
@@ -36,11 +36,11 @@ def _fwd_kernel_destindex_copy_kv(
 
     dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
 
-    kv_nope_ptrs = KV_nope + cur_index * stride_kv_nope_bs + stride_kv_nope_d * offs_d_nope[None, :]
-    kv_rope_ptrs = KV_rope + cur_index * stride_kv_rope_bs + stride_kv_rope_d * offs_d_rope[None, :]
+    kv_nope_ptrs = KV_nope + cur_index * stride_kv_nope_bs + stride_kv_nope_d * offs_d_nope
+    kv_rope_ptrs = KV_rope + cur_index * stride_kv_rope_bs + stride_kv_rope_d * offs_d_rope
 
-    o_nope_ptrs = O_nope + dest_index * stride_o_nope_bs + stride_o_nope_d * offs_d_nope[None, :]
-    o_rope_ptrs = O_rope + dest_index * stride_o_rope_bs + stride_o_rope_d * offs_d_rope[None, :]
+    o_nope_ptrs = O_nope + dest_index * stride_o_nope_bs + stride_o_nope_d * offs_d_nope
+    o_rope_ptrs = O_rope + dest_index * stride_o_rope_bs + stride_o_rope_d * offs_d_rope
 
     kv_nope = tl.load(kv_nope_ptrs)
     kv_rope = tl.load(kv_rope_ptrs)
@@ -60,6 +60,9 @@ def destindex_copy_kv(KV_nope, KV_rope, DestLoc, O_nope, O_rope):
     assert KV_nope.shape[2] == O_nope.shape[2]
     assert KV_rope.shape[1] == O_rope.shape[1]
     assert KV_rope.shape[2] == O_rope.shape[2]
+    assert triton.next_power_of_2(kv_nope_head_dim) == kv_nope_head_dim
+    assert triton.next_power_of_2(kv_rope_head_dim) == kv_rope_head_dim
+
     grid = (seq_len,)
     num_warps = 1
 
diff --git a/lightllm/common/basemodel/triton_kernel/kv_copy/ppl_int4kv_copy_kv.py b/lightllm/common/basemodel/triton_kernel/kv_copy/ppl_int4kv_copy_kv.py
new file mode 100644
index 0000000000..53d1256ec7
--- /dev/null
+++ b/lightllm/common/basemodel/triton_kernel/kv_copy/ppl_int4kv_copy_kv.py
@@ -0,0 +1,374 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _fwd_kernel_destindex_copy_quantize_int4_kv(
+    K,
+    Dest_loc,
+    Out,
+    Out_scale,
+    stride_k_bs,
+    stride_k_h,
+    stride_k_g,
+    stride_k_d,
+    stride_o_bs,
+    stride_o_h,
+    stride_o_g,
+    stride_o_d,
+    stride_os_bs,
+    stride_os_h,
+    stride_os_g,
+    group_count,
+    token_num,
+    HEAD_NUM: tl.constexpr,
+    BLOCK_GROUP_COUNT: tl.constexpr,
+    BLOCK_GROUP_DIM: tl.constexpr,
+):
+    start_index = tl.program_id(0)
+
+    for cur_index in range(start_index, token_num, step=tl.num_programs(axis=0)):
+        offs_g = tl.arange(0, BLOCK_GROUP_COUNT) % group_count
+        offs_d = tl.arange(0, BLOCK_GROUP_DIM // 2)
+
+        dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
+
+        for cur_head in tl.static_range(HEAD_NUM, step=1):
+            src_data_0 = tl.load(
+                K
+                + cur_index * stride_k_bs
+                + cur_head * stride_k_h
+                + offs_g[:, None] * stride_k_g
+                + offs_d[None, :] * 2,
+            )
+            src_data_1 = tl.load(
+                K
+                + cur_index * stride_k_bs
+                + cur_head * stride_k_h
+                + offs_g[:, None] * stride_k_g
+                + offs_d[None, :] * 2
+                + 1,
+            )
+
+            abs_data_0 = tl.abs(src_data_0)
+            abs_data_1 = tl.abs(src_data_1)
+
+            data_scale = (tl.maximum(tl.max(abs_data_0, axis=1), tl.max(abs_data_1, axis=1)) / 7.0).to(
+                Out_scale.dtype.element_ty
+            )
+            q_src_data_0 = (src_data_0 / data_scale[:, None]).to(tl.int8)
+            q_src_data_0 = tl.where(q_src_data_0 > 7, 7, q_src_data_0)
+            q_src_data_0 = tl.where(q_src_data_0 < -7, -7, q_src_data_0)
+            q_src_data_0 += 7
+            q_src_data_0 = q_src_data_0.to(tl.uint8, bitcast=True)
+
+            q_src_data_1 = (src_data_1 / data_scale[:, None]).to(tl.int8)
+            q_src_data_1 = tl.where(q_src_data_1 > 7, 7, q_src_data_1)
+            q_src_data_1 = tl.where(q_src_data_1 < -7, -7, q_src_data_1)
+            q_src_data_1 += 7
+            q_src_data_1 = q_src_data_1.to(tl.uint8, bitcast=True)
+
+            low_4 = q_src_data_0 & 0xF
+            high_4 = (q_src_data_1 & 0xF) << 4
+
+            out_data = (low_4 | high_4).to(Out.dtype.element_ty, bitcast=True)
+
+            o_ptrs = (
+                Out + dest_index * stride_o_bs + cur_head * stride_o_h + offs_g[:, None] * stride_o_g + offs_d[None, :]
+            )
+            os_ptrs = Out_scale + dest_index * stride_os_bs + cur_head * stride_os_h + offs_g
+            tl.store(o_ptrs, out_data)
+            tl.store(os_ptrs, data_scale)
+    return
+
+
+@torch.no_grad()
+def destindex_copy_int4kv(
+    KV: torch.Tensor,
+    DestLoc: torch.Tensor,
+    KV_buffer: torch.Tensor,
+    KV_scale_buffer: torch.Tensor,
+    quant_group_size: int,
+):
+    head_num = KV.shape[1]
+    head_dim = KV.shape[2]
+
+    assert head_dim % quant_group_size == 0, "error head dim, can not been supported to copy quant kv"
+
+    group_count = head_dim // quant_group_size
+    group_dim = quant_group_size
+
+    assert triton.next_power_of_2(group_dim) == group_dim
+
+    KV = KV.view((KV.shape[0], head_num, group_count, group_dim))
+    KV_buffer = KV_buffer.view(
+        KV_buffer.shape[0], KV_buffer.shape[1], group_count, group_dim // 2
+    )  # OUt 是 int8 类型， 两个int4组一个int8，所以 group_dim // 2
+    KV_scale_buffer = KV_scale_buffer.view(KV_scale_buffer.shape[0], KV_scale_buffer.shape[1], group_count)
+    if len(DestLoc) < 1024:
+        grid = (len(DestLoc),)
+    else:
+        grid = (1024,)
+
+    _fwd_kernel_destindex_copy_quantize_int4_kv[grid](
+        K=KV,
+        Dest_loc=DestLoc,
+        Out=KV_buffer,
+        Out_scale=KV_scale_buffer,
+        stride_k_bs=KV.stride(0),
+        stride_k_h=KV.stride(1),
+        stride_k_g=KV.stride(2),
+        stride_k_d=KV.stride(3),
+        stride_o_bs=KV_buffer.stride(0),
+        stride_o_h=KV_buffer.stride(1),
+        stride_o_g=KV_buffer.stride(2),
+        stride_o_d=KV_buffer.stride(3),
+        stride_os_bs=KV_scale_buffer.stride(0),
+        stride_os_h=KV_scale_buffer.stride(1),
+        stride_os_g=KV_scale_buffer.stride(2),
+        group_count=group_count,
+        token_num=len(DestLoc),
+        HEAD_NUM=head_num,
+        BLOCK_GROUP_COUNT=triton.next_power_of_2(group_count),
+        BLOCK_GROUP_DIM=group_dim,
+        num_warps=4,
+        num_stages=1,
+    )
+    return
+
+
+@triton.jit
+def int4_to_float(k_int8, offs_d):
+    k_int8 = k_int8.to(tl.uint8, bitcast=True)
+    k_high = (k_int8 & 0xF0) >> 4
+    k_low = k_int8 & 0x0F
+    k_high = k_high.to(tl.int8, bitcast=True)
+    k_low = k_low.to(tl.int8, bitcast=True)
+    k_high -= 7
+    k_low -= 7
+
+    k_int4 = tl.where(
+        offs_d[None, None, :] % 2 == 0,
+        k_low,
+        k_high,
+    )
+    return k_int4
+
+
+@triton.jit
+def _fwd_dequantize_int4kv(
+    k,
+    k_ss,
+    k_sh,
+    k_sg,
+    k_sd,
+    k_scale,
+    k_scale_ss,
+    k_scale_sh,
+    k_scale_sg,
+    k_scale_sd,
+    v,
+    v_ss,
+    v_sh,
+    v_sg,
+    v_sd,
+    v_scale,
+    v_scale_ss,
+    v_scale_sh,
+    v_scale_sg,
+    v_scale_sd,
+    req_to_token_indexs,
+    stride_req_to_tokens_b,
+    stride_req_to_tokens_s,
+    b_seq_len,
+    b_req_idx,
+    b_kv_start_loc,
+    k_out,
+    k_out_ss,
+    k_out_sh,
+    k_out_sg,
+    k_out_sd,
+    v_out,
+    v_out_ss,
+    v_out_sh,
+    v_out_sg,
+    v_out_sd,
+    k_head_num,
+    v_head_num,
+    group_count,
+    group_dim,
+    SEQ_BLOCK_SIZE: tl.constexpr,
+    GROUP_COUNT_BLOCK_SIZE: tl.constexpr,
+    BLOCK_GROUP_DIM: tl.constexpr,
+):
+    start_block_index = tl.program_id(0)
+    cur_batch = tl.program_id(1)
+    cur_batch_req_idx = tl.load(b_req_idx + cur_batch)
+    cur_seq_len = tl.load(b_seq_len + cur_batch)
+    if start_block_index * SEQ_BLOCK_SIZE >= cur_seq_len:
+        return
+
+    out_start_loc = tl.load(b_kv_start_loc + cur_batch)
+
+    offs_kv_loc = (start_block_index * SEQ_BLOCK_SIZE + tl.arange(0, SEQ_BLOCK_SIZE)) % cur_seq_len
+    kv_loc = tl.load(req_to_token_indexs + cur_batch_req_idx * stride_req_to_tokens_b + offs_kv_loc).to(tl.int64)
+
+    offs_d = tl.arange(0, BLOCK_GROUP_DIM)
+    offs_scale_d = tl.arange(0, 1)
+    group_offs = tl.arange(0, GROUP_COUNT_BLOCK_SIZE) % group_count
+
+    for k_head_index in tl.range(0, k_head_num, step=1, num_stages=3):
+        k_int8 = tl.load(
+            k
+            + kv_loc[:, None, None] * k_ss
+            + k_head_index * k_sh
+            + group_offs[None, :, None] * k_sg
+            + offs_d[None, None, :] // 2
+        )
+        k_int4 = int4_to_float(k_int8, offs_d)
+
+        k_scale_data = tl.load(
+            k_scale
+            + kv_loc[:, None, None] * k_scale_ss
+            + k_head_index * k_scale_sh
+            + group_offs[None, :, None] * k_scale_sg
+            + offs_scale_d[None, None, :]
+        )
+        k_out_data = k_int4.to(k_out.dtype.element_ty) * k_scale_data
+        tl.store(
+            k_out
+            + (out_start_loc + offs_kv_loc[:, None, None]) * k_out_ss
+            + k_head_index * k_out_sh
+            + group_offs[None, :, None] * k_out_sg
+            + offs_d[None, None, :],
+            k_out_data,
+        )
+
+    for v_head_index in tl.range(0, v_head_num, step=1, num_stages=3):
+        v_int8 = tl.load(
+            v
+            + kv_loc[:, None, None] * v_ss
+            + v_head_index * v_sh
+            + group_offs[None, :, None] * v_sg
+            + offs_d[None, None, :] // 2
+        )
+        v_int4 = int4_to_float(v_int8, offs_d)
+        v_scale_data = tl.load(
+            v_scale
+            + kv_loc[:, None, None] * v_scale_ss
+            + v_head_index * v_scale_sh
+            + group_offs[None, :, None] * v_scale_sg
+            + offs_scale_d[None, None, :]
+        )
+        v_out_data = v_int4.to(v_out.dtype.element_ty) * v_scale_data
+        tl.store(
+            v_out
+            + (out_start_loc + offs_kv_loc[:, None, None]) * v_out_ss
+            + v_head_index * v_out_sh
+            + group_offs[None, :, None] * v_out_sg
+            + offs_d[None, None, :],
+            v_out_data,
+        )
+    return
+
+
+@torch.no_grad()
+def dequantize_int4kv(
+    k: torch.Tensor,
+    k_scale: torch.Tensor,
+    v: torch.Tensor,
+    v_scale: torch.Tensor,
+    req_to_token_indexs: torch.Tensor,
+    b_seq_len: torch.Tensor,
+    b_req_idx: torch.Tensor,
+    b_kv_start_loc: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+    max_len_in_batch: int,
+    quant_group_size: int,
+):
+    batch_size = b_seq_len.shape[0]
+    k_head_num = k.shape[1]
+    k_head_dim = k.shape[2] * 2
+    v_head_num = v.shape[1]
+    v_head_dim = v.shape[2] * 2
+    assert k_head_dim % quant_group_size == 0, "error head dim, can not been supported to copy quant kv"
+    assert v_head_dim % quant_group_size == 0, "error head dim, can not been supported to copy quant kv"
+    assert k_head_dim == v_head_dim, "error head dim, can not been supported to copy quant kv"
+    assert k_head_dim // v_scale.shape[2] == quant_group_size, "error head dim, can not been supported to copy quant kv"
+    assert k_head_dim in [64, 128, 256]
+
+    group_count = k_head_dim // quant_group_size
+    group_dim = quant_group_size
+
+    assert triton.next_power_of_2(group_dim) == group_dim
+
+    k = k.view((k.shape[0], k.shape[1], group_count, group_dim // 2))  # int4kv 以 int8 存储的
+    v = v.view((v.shape[0], v.shape[1], group_count, group_dim // 2))
+    k_scale = k_scale.view((k_scale.shape[0], k_scale.shape[1], group_count, 1))
+    v_scale = v_scale.view((v_scale.shape[0], v_scale.shape[1], group_count, 1))
+
+    # 使拆分的grid 具有足够的并行度
+    SEQ_BLOCK_SIZE = 128
+    while triton.cdiv(max_len_in_batch, SEQ_BLOCK_SIZE) * batch_size < 512:
+        SEQ_BLOCK_SIZE = SEQ_BLOCK_SIZE // 2
+        if SEQ_BLOCK_SIZE <= 1:
+            break
+
+    if SEQ_BLOCK_SIZE <= 1:
+        SEQ_BLOCK_SIZE = 8
+
+    grid = (triton.cdiv(max_len_in_batch, SEQ_BLOCK_SIZE), batch_size)
+    num_warps = 4
+    k_out = k_out.view((k_out.shape[0], k_out.shape[1], group_count, group_dim))
+    v_out = v_out.view((v_out.shape[0], v_out.shape[1], group_count, group_dim))
+
+    _fwd_dequantize_int4kv[grid](
+        k=k,
+        k_ss=k.stride(0),
+        k_sh=k.stride(1),
+        k_sg=k.stride(2),
+        k_sd=k.stride(3),
+        k_scale=k_scale,
+        k_scale_ss=k_scale.stride(0),
+        k_scale_sh=k_scale.stride(1),
+        k_scale_sg=k_scale.stride(2),
+        k_scale_sd=k_scale.stride(2),
+        v=v,
+        v_ss=v.stride(0),
+        v_sh=v.stride(1),
+        v_sg=v.stride(2),
+        v_sd=v.stride(3),
+        v_scale=v_scale,
+        v_scale_ss=v_scale.stride(0),
+        v_scale_sh=v_scale.stride(1),
+        v_scale_sg=v_scale.stride(2),
+        v_scale_sd=v_scale.stride(3),
+        req_to_token_indexs=req_to_token_indexs,
+        stride_req_to_tokens_b=req_to_token_indexs.stride(0),
+        stride_req_to_tokens_s=req_to_token_indexs.stride(1),
+        b_seq_len=b_seq_len,
+        b_req_idx=b_req_idx,
+        b_kv_start_loc=b_kv_start_loc,
+        k_out=k_out,
+        k_out_ss=k_out.stride(0),
+        k_out_sh=k_out.stride(1),
+        k_out_sg=k_out.stride(2),
+        k_out_sd=k_out.stride(3),
+        v_out=v_out,
+        v_out_ss=v_out.stride(0),
+        v_out_sh=v_out.stride(1),
+        v_out_sg=v_out.stride(2),
+        v_out_sd=v_out.stride(3),
+        k_head_num=k_head_num,
+        v_head_num=v_head_num,
+        group_count=group_count,
+        group_dim=group_dim,
+        SEQ_BLOCK_SIZE=SEQ_BLOCK_SIZE,
+        GROUP_COUNT_BLOCK_SIZE=triton.next_power_of_2(group_count),
+        BLOCK_GROUP_DIM=group_dim,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+    return
diff --git a/lightllm/common/basemodel/triton_kernel/kv_copy/ppl_int8kv_copy_kv.py b/lightllm/common/basemodel/triton_kernel/kv_copy/ppl_int8kv_copy_kv.py
new file mode 100644
index 0000000000..e5ee5cb8b8
--- /dev/null
+++ b/lightllm/common/basemodel/triton_kernel/kv_copy/ppl_int8kv_copy_kv.py
@@ -0,0 +1,330 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _fwd_kernel_destindex_copy_quantize_kv(
+    K,
+    Dest_loc,
+    Out,
+    Out_scale,
+    stride_k_bs,
+    stride_k_h,
+    stride_k_g,
+    stride_k_d,
+    stride_o_bs,
+    stride_o_h,
+    stride_o_g,
+    stride_o_d,
+    stride_os_bs,
+    stride_os_h,
+    stride_os_g,
+    group_size,
+    BLOCK_GROUP_NUM: tl.constexpr,
+    BLOCK_GROUP_DIM: tl.constexpr,
+):
+    cur_index = tl.program_id(0)
+    cur_head = tl.program_id(1)
+
+    offs_g = tl.arange(0, BLOCK_GROUP_NUM)
+    offs_d = tl.arange(0, BLOCK_GROUP_DIM)
+
+    dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
+
+    src_data = tl.load(
+        K + cur_index * stride_k_bs + cur_head * stride_k_h + offs_g[:, None] * stride_k_g + offs_d[None, :],
+        mask=offs_g[:, None] < group_size,
+        other=0.0,
+    )
+    abs_data = tl.abs(src_data)
+    data_scale = (tl.max(abs_data, axis=1) / 127.0).to(Out_scale.dtype.element_ty)
+    q_src_data = (src_data / data_scale[:, None]).to(tl.int8)
+
+    o_ptrs = Out + dest_index * stride_o_bs + cur_head * stride_o_h + offs_g[:, None] * stride_o_g + offs_d[None, :]
+    os_ptrs = Out_scale + dest_index * stride_os_bs + cur_head * stride_os_h + offs_g
+    tl.store(o_ptrs, q_src_data, mask=offs_g[:, None] < group_size)
+    tl.store(os_ptrs, data_scale, mask=offs_g < group_size)
+    return
+
+
+@torch.no_grad()
+def destindex_copy_quantize_kv(K, DestLoc, Out, Out_scale, quant_group_dim):
+    seq_len = DestLoc.shape[0]
+    head_num = K.shape[1]
+    head_dim = K.shape[2]
+    assert triton.next_power_of_2(quant_group_dim) == quant_group_dim, "error quant group dim"
+
+    assert head_dim % quant_group_dim == 0, "error head dim, can not been supported to copy quant kv"
+    grid = (seq_len, head_num)
+    num_warps = 1
+
+    group_size = head_dim // quant_group_dim
+    group_dim = quant_group_dim
+
+    K = K.view((K.shape[0], K.shape[1], group_size, group_dim))
+    Out = Out.view(Out.shape[0], Out.shape[1], group_size, group_dim)
+
+    _fwd_kernel_destindex_copy_quantize_kv[grid](
+        K,
+        DestLoc,
+        Out,
+        Out_scale,
+        K.stride(0),
+        K.stride(1),
+        K.stride(2),
+        K.stride(3),
+        Out.stride(0),
+        Out.stride(1),
+        Out.stride(2),
+        Out.stride(3),
+        Out_scale.stride(0),
+        Out_scale.stride(1),
+        Out_scale.stride(2),
+        group_size,
+        BLOCK_GROUP_NUM=triton.next_power_of_2(group_size),
+        BLOCK_GROUP_DIM=group_dim,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+    return
+
+
+@triton.jit
+def _fwd_dequantize_int8kv(
+    k,
+    k_ss,
+    k_sh,
+    k_sg,
+    k_sd,
+    k_scale,
+    k_scale_ss,
+    k_scale_sh,
+    k_scale_sg,
+    k_scale_sd,
+    v,
+    v_ss,
+    v_sh,
+    v_sg,
+    v_sd,
+    v_scale,
+    v_scale_ss,
+    v_scale_sh,
+    v_scale_sg,
+    v_scale_sd,
+    req_to_token_indexs,
+    stride_req_to_tokens_b,
+    stride_req_to_tokens_s,
+    b_seq_len,
+    b_req_idx,
+    b_kv_start_loc,
+    k_out,
+    k_out_ss,
+    k_out_sh,
+    k_out_sg,
+    k_out_sd,
+    v_out,
+    v_out_ss,
+    v_out_sh,
+    v_out_sg,
+    v_out_sd,
+    k_head_num,
+    v_head_num,
+    group_count,
+    group_dim,
+    SEQ_BLOCK_SIZE: tl.constexpr,
+    GROUP_COUNT_BLOCK_SIZE: tl.constexpr,
+    BLOCK_GROUP_DIM: tl.constexpr,
+):
+    start_block_index = tl.program_id(0)
+    cur_batch = tl.program_id(1)
+    cur_batch_req_idx = tl.load(b_req_idx + cur_batch)
+    cur_seq_len = tl.load(b_seq_len + cur_batch)
+    if start_block_index * SEQ_BLOCK_SIZE >= cur_seq_len:
+        return
+
+    out_start_loc = tl.load(b_kv_start_loc + cur_batch)
+
+    offs_kv_loc = (start_block_index * SEQ_BLOCK_SIZE + tl.arange(0, SEQ_BLOCK_SIZE)) % cur_seq_len
+    kv_loc = tl.load(req_to_token_indexs + cur_batch_req_idx * stride_req_to_tokens_b + offs_kv_loc).to(tl.int64)
+
+    offs_d = tl.arange(0, BLOCK_GROUP_DIM) % group_dim
+    offs_scale_d = tl.arange(0, 1)
+    group_offs = tl.arange(0, GROUP_COUNT_BLOCK_SIZE) % group_count
+
+    for k_head_index in tl.range(0, k_head_num, step=1, num_stages=3):
+        k_int8 = tl.load(
+            k
+            + kv_loc[:, None, None] * k_ss
+            + k_head_index * k_sh
+            + group_offs[None, :, None] * k_sg
+            + offs_d[None, None, :]
+        )
+        k_scale_data = tl.load(
+            k_scale
+            + kv_loc[:, None, None] * k_scale_ss
+            + k_head_index * k_scale_sh
+            + group_offs[None, :, None] * k_scale_sg
+            + offs_scale_d[None, None, :]
+        )
+        k_out_data = k_int8.to(k_out.dtype.element_ty) * k_scale_data
+        tl.store(
+            k_out
+            + (out_start_loc + offs_kv_loc[:, None, None]) * k_out_ss
+            + k_head_index * k_out_sh
+            + group_offs[None, :, None] * k_out_sg
+            + offs_d[None, None, :],
+            k_out_data,
+        )
+
+    for v_head_index in tl.range(0, v_head_num, step=1, num_stages=3):
+        v_int8 = tl.load(
+            v
+            + kv_loc[:, None, None] * v_ss
+            + v_head_index * v_sh
+            + group_offs[None, :, None] * v_sg
+            + offs_d[None, None, :]
+        )
+        v_scale_data = tl.load(
+            v_scale
+            + kv_loc[:, None, None] * v_scale_ss
+            + v_head_index * v_scale_sh
+            + group_offs[None, :, None] * v_scale_sg
+            + offs_scale_d[None, None, :]
+        )
+        v_out_data = v_int8.to(v_out.dtype.element_ty) * v_scale_data
+        tl.store(
+            v_out
+            + (out_start_loc + offs_kv_loc[:, None, None]) * v_out_ss
+            + v_head_index * v_out_sh
+            + group_offs[None, :, None] * v_out_sg
+            + offs_d[None, None, :],
+            v_out_data,
+        )
+    return
+
+
+@torch.no_grad()
+def dequantize_int8kv(
+    k: torch.Tensor,
+    k_scale: torch.Tensor,
+    v: torch.Tensor,
+    v_scale: torch.Tensor,
+    req_to_token_indexs: torch.Tensor,
+    b_seq_len: torch.Tensor,
+    b_req_idx: torch.Tensor,
+    b_kv_start_loc: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+    max_len_in_batch: int,
+    quant_group_size: int,
+):
+    batch_size = b_seq_len.shape[0]
+    k_head_num = k.shape[1]
+    k_head_dim = k.shape[2]
+    v_head_num = v.shape[1]
+    v_head_dim = v.shape[2]
+    assert k_head_dim % quant_group_size == 0, "error head dim, can not been supported to copy quant kv"
+    assert v_head_dim % quant_group_size == 0, "error head dim, can not been supported to copy quant kv"
+    assert k_head_dim == v_head_dim, "error head dim, can not been supported to copy quant kv"
+    assert k_head_dim // v_scale.shape[2] == quant_group_size, "error head dim, can not been supported to copy quant kv"
+    assert k_head_dim in [64, 128, 256]
+
+    group_count = k_head_dim // quant_group_size
+    group_dim = quant_group_size
+
+    k = k.view((k.shape[0], k.shape[1], group_count, group_dim))
+    v = v.view((v.shape[0], v.shape[1], group_count, group_dim))
+    k_scale = k_scale.view((k_scale.shape[0], k_scale.shape[1], group_count, 1))
+    v_scale = v_scale.view((v_scale.shape[0], v_scale.shape[1], group_count, 1))
+
+    # 使拆分的grid 具有足够的并行度
+    SEQ_BLOCK_SIZE = 128
+    while triton.cdiv(max_len_in_batch, SEQ_BLOCK_SIZE) * batch_size < 512:
+        SEQ_BLOCK_SIZE = SEQ_BLOCK_SIZE // 2
+        if SEQ_BLOCK_SIZE <= 1:
+            break
+
+    if SEQ_BLOCK_SIZE <= 1:
+        SEQ_BLOCK_SIZE = 8
+
+    grid = (triton.cdiv(max_len_in_batch, SEQ_BLOCK_SIZE), batch_size)
+    num_warps = 4
+    k_out = k_out.view((k_out.shape[0], k_out.shape[1], group_count, group_dim))
+    v_out = v_out.view((v_out.shape[0], v_out.shape[1], group_count, group_dim))
+
+    _fwd_dequantize_int8kv[grid](
+        k=k,
+        k_ss=k.stride(0),
+        k_sh=k.stride(1),
+        k_sg=k.stride(2),
+        k_sd=k.stride(3),
+        k_scale=k_scale,
+        k_scale_ss=k_scale.stride(0),
+        k_scale_sh=k_scale.stride(1),
+        k_scale_sg=k_scale.stride(2),
+        k_scale_sd=k_scale.stride(2),
+        v=v,
+        v_ss=v.stride(0),
+        v_sh=v.stride(1),
+        v_sg=v.stride(2),
+        v_sd=v.stride(3),
+        v_scale=v_scale,
+        v_scale_ss=v_scale.stride(0),
+        v_scale_sh=v_scale.stride(1),
+        v_scale_sg=v_scale.stride(2),
+        v_scale_sd=v_scale.stride(3),
+        req_to_token_indexs=req_to_token_indexs,
+        stride_req_to_tokens_b=req_to_token_indexs.stride(0),
+        stride_req_to_tokens_s=req_to_token_indexs.stride(1),
+        b_seq_len=b_seq_len,
+        b_req_idx=b_req_idx,
+        b_kv_start_loc=b_kv_start_loc,
+        k_out=k_out,
+        k_out_ss=k_out.stride(0),
+        k_out_sh=k_out.stride(1),
+        k_out_sg=k_out.stride(2),
+        k_out_sd=k_out.stride(3),
+        v_out=v_out,
+        v_out_ss=v_out.stride(0),
+        v_out_sh=v_out.stride(1),
+        v_out_sg=v_out.stride(2),
+        v_out_sd=v_out.stride(3),
+        k_head_num=k_head_num,
+        v_head_num=v_head_num,
+        group_count=group_count,
+        group_dim=group_dim,
+        SEQ_BLOCK_SIZE=SEQ_BLOCK_SIZE,
+        GROUP_COUNT_BLOCK_SIZE=triton.next_power_of_2(group_count),
+        BLOCK_GROUP_DIM=triton.next_power_of_2(group_dim),
+        num_warps=num_warps,
+        num_stages=1,
+    )
+    return
+
+
+def test2():
+    import time
+
+    B, N_CTX, H, D = 1, 3, 12, 128
+    src = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda()
+    dest_loc = torch.arange(0, B * N_CTX, dtype=torch.int32).cuda()
+    value_dest = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda().to(torch.int8)
+    scale_dest = torch.randn((B * N_CTX, H, D // 8), dtype=torch.float16).cuda()
+
+    for _ in range(10):
+        destindex_copy_quantize_kv(src, dest_loc, value_dest, scale_dest)
+    torch.cuda.synchronize()
+    t1 = time.time()
+    for _ in range(1000):
+        destindex_copy_quantize_kv(src, dest_loc, value_dest, scale_dest)
+    torch.cuda.synchronize()
+    t2 = time.time()
+
+    print("Time cost ", t2 - t1)
+    value_dest = value_dest.view((B * N_CTX, H, D // 8, 8))
+    scale_dest = scale_dest.view((B * N_CTX, H, D // 8, 1))
+    print("max ", torch.max(torch.abs((value_dest * scale_dest).view(B * N_CTX, H, D) - src)))
+    print("mean ", torch.mean(torch.abs((value_dest * scale_dest).view(B * N_CTX, H, D) - src)))
+    cos = torch.nn.CosineSimilarity(0)
+    print("cos ", cos(src.flatten().to(torch.float32), (value_dest * scale_dest).flatten().to(torch.float32)))
diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/__init__.py b/lightllm/common/basemodel/triton_kernel/mla_att/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/__init__.py b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/__init__.py
new file mode 100644
index 0000000000..fb0609a401
--- /dev/null
+++ b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/__init__.py
@@ -0,0 +1 @@
+from .gqa_flash_decoding import gqa_token_decode_attention_flash_decoding
diff --git a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py
similarity index 92%
rename from lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding.py
rename to lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py
index 256dfce5af..28839b5f59 100644
--- a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding.py
+++ b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py
@@ -12,28 +12,21 @@
 
 
 def gqa_token_decode_attention_flash_decoding(
-    q_nope,
-    q_rope,
-    kv_nope,
-    kv_rope,
-    infer_state,
-    q_head_num,
-    kv_lora_rank,
-    q_rope_dim,
-    qk_nope_head_dim,
-    softmax_scale,
-    out=None,
-    alloc_tensor_func=torch.empty,
-    **run_config
+    q_nope, q_rope, kv_nope, kv_rope, infer_state, softmax_scale, out=None, alloc_tensor_func=torch.empty, **run_config
 ):
     batch_size = infer_state.batch_size
-    max_len_in_batch = infer_state.max_len_in_batch
+    max_kv_seq_len = infer_state.max_kv_seq_len
+
+    q_head_num, kv_lora_rank = q_nope.shape[1], q_nope.shape[2]
+    q_rope_dim = q_rope.shape[2]
+    assert q_rope_dim == 64
+
     calcu_shape1 = (batch_size, q_head_num, kv_lora_rank)
     calcu_shape2 = (batch_size, q_head_num, q_rope_dim)
 
     if not run_config:
         if torch.cuda.is_current_stream_capturing():
-            avg_seq_len_in_batch = max_len_in_batch
+            avg_seq_len_in_batch = max_kv_seq_len
         else:
             avg_seq_len_in_batch = infer_state.total_token_num // batch_size
 
diff --git a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_config.py b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_config.py
similarity index 100%
rename from lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_config.py
rename to lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_config.py
diff --git a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py
similarity index 100%
rename from lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py
rename to lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py
diff --git a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py
similarity index 100%
rename from lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py
rename to lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py
diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py b/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py
new file mode 100644
index 0000000000..5725bed2e7
--- /dev/null
+++ b/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py
@@ -0,0 +1 @@
+from .context_flashattention_nopad_with_v import context_attention_fwd_with_v
diff --git a/lightllm/models/deepseek2/triton_kernel/context_flashattention_nopad_with_v.py b/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/context_flashattention_nopad_with_v.py
similarity index 100%
rename from lightllm/models/deepseek2/triton_kernel/context_flashattention_nopad_with_v.py
rename to lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/context_flashattention_nopad_with_v.py
diff --git a/lightllm/common/basemodel/triton_kernel/norm/__init__.py b/lightllm/common/basemodel/triton_kernel/norm/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/common/basemodel/triton_kernel/layernorm.py b/lightllm/common/basemodel/triton_kernel/norm/layernorm.py
similarity index 100%
rename from lightllm/common/basemodel/triton_kernel/layernorm.py
rename to lightllm/common/basemodel/triton_kernel/norm/layernorm.py
diff --git a/lightllm/models/qwen3/triton_kernel/qk_norm.py b/lightllm/common/basemodel/triton_kernel/norm/qk_norm.py
similarity index 100%
rename from lightllm/models/qwen3/triton_kernel/qk_norm.py
rename to lightllm/common/basemodel/triton_kernel/norm/qk_norm.py
diff --git a/lightllm/common/basemodel/triton_kernel/rmsnorm.py b/lightllm/common/basemodel/triton_kernel/norm/rmsnorm.py
similarity index 100%
rename from lightllm/common/basemodel/triton_kernel/rmsnorm.py
rename to lightllm/common/basemodel/triton_kernel/norm/rmsnorm.py
diff --git a/lightllm/common/basemodel/triton_kernel/quantization/__init__.py b/lightllm/common/basemodel/triton_kernel/quantization/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/common/basemodel/triton_kernel/bmm_scaled_fp8.py b/lightllm/common/basemodel/triton_kernel/quantization/bmm_scaled_fp8.py
similarity index 100%
rename from lightllm/common/basemodel/triton_kernel/bmm_scaled_fp8.py
rename to lightllm/common/basemodel/triton_kernel/quantization/bmm_scaled_fp8.py
diff --git a/lightllm/common/quantization/triton_quant/fp8/fp8act_quant_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py
similarity index 100%
rename from lightllm/common/quantization/triton_quant/fp8/fp8act_quant_kernel.py
rename to lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py
diff --git a/lightllm/common/quantization/triton_quant/fp8/fp8w8a8_block_gemm_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_gemm_kernel.py
similarity index 100%
rename from lightllm/common/quantization/triton_quant/fp8/fp8w8a8_block_gemm_kernel.py
rename to lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_gemm_kernel.py
diff --git a/lightllm/common/quantization/triton_quant/fp8/fp8w8a8_block_quant_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py
similarity index 100%
rename from lightllm/common/quantization/triton_quant/fp8/fp8w8a8_block_quant_kernel.py
rename to lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py
diff --git a/lightllm/common/basemodel/triton_kernel/q_per_head_fp8_quant.py b/lightllm/common/basemodel/triton_kernel/quantization/q_per_head_fp8_quant.py
similarity index 100%
rename from lightllm/common/basemodel/triton_kernel/q_per_head_fp8_quant.py
rename to lightllm/common/basemodel/triton_kernel/quantization/q_per_head_fp8_quant.py
diff --git a/lightllm/common/quantization/triton_quant/fp8/fp8w8a8_scaled_mm_per_token_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/scaled_mm_per_token_kernel.py
similarity index 93%
rename from lightllm/common/quantization/triton_quant/fp8/fp8w8a8_scaled_mm_per_token_kernel.py
rename to lightllm/common/basemodel/triton_kernel/quantization/scaled_mm_per_token_kernel.py
index 7c76e82c9e..f14e8b2833 100644
--- a/lightllm/common/quantization/triton_quant/fp8/fp8w8a8_scaled_mm_per_token_kernel.py
+++ b/lightllm/common/basemodel/triton_kernel/quantization/scaled_mm_per_token_kernel.py
@@ -11,8 +11,8 @@
 from lightllm.utils.device_utils import triton_support_tensor_descriptor, is_5090_gpu
 
 
-class Fp8ScaledMMKernelConfig(KernelConfigs):
-    kernel_name: str = "fp8_scaled_mm_per_token"
+class ScaledMMKernelConfig(KernelConfigs):
+    kernel_name: str = "scaled_mm_per_token"
 
     @classmethod
     @lru_cache(maxsize=200)
@@ -105,6 +105,7 @@ def _scaled_mm_per_token(
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
     GROUP_M: tl.constexpr,
+    ACC_DTYPE: tl.constexpr,
 ):
     pid = tl.program_id(0)
     m_block_num = tl.cdiv(M, BLOCK_M)
@@ -134,7 +135,7 @@ def _scaled_mm_per_token(
     a_s = tl.load(Ascale_ptrs)
     b_s = tl.load(Bscale_ptrs)
 
-    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
 
     for k in range(0, tl.cdiv(K, BLOCK_K)):
         if USE_TMA:
@@ -155,6 +156,7 @@ def _scaled_mm_per_token(
             a_ptrs += BLOCK_K * stride_ak
             b_ptrs += BLOCK_K * stride_bk
 
+    acc = acc.to(tl.float32)
     acc = acc * a_s[:, None] * b_s[None, :]
 
     acc = acc.to(out.dtype.element_ty)
@@ -206,13 +208,13 @@ def _get_static_key(A, B, out_dtype):
 
 
 @autotune(
-    kernel_name="fp8_scaled_mm_per_token:v3",
+    kernel_name="scaled_mm_per_token:v1",
     configs_gen_func=get_test_configs,
     static_key_func=_get_static_key,
     run_key_func=lambda A: A.shape[0],
     mutates_args=["out"],
 )
-def fp8_scaled_mm_per_token(
+def scaled_mm_per_token(
     A: torch.Tensor,
     B: torch.Tensor,
     Ascale: torch.Tensor,
@@ -221,7 +223,7 @@ def fp8_scaled_mm_per_token(
     out: torch.Tensor,
     run_config=None,
 ) -> torch.Tensor:
-    """w8a8fp8 per-token quantization mm.
+    """w8a8 per-token quantization mm (supports fp8 and int8).
 
     Args:
         A: Matrix A with shape of [M, K].
@@ -239,7 +241,7 @@ def fp8_scaled_mm_per_token(
     M, K = A.shape
     _, N = B.shape
     if not run_config:
-        run_config = Fp8ScaledMMKernelConfig.try_to_get_best_config(M=M, N=N, K=K, out_dtype=out_dtype)
+        run_config = ScaledMMKernelConfig.try_to_get_best_config(M=M, N=N, K=K, out_dtype=out_dtype)
     NEED_N_MASK = N % run_config["BLOCK_N"] != 0
     NEED_K_MASK = K % run_config["BLOCK_K"] != 0
     grid = (triton.cdiv(M, run_config["BLOCK_M"]) * triton.cdiv(N, run_config["BLOCK_N"]),)
@@ -283,6 +285,8 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]):
         B_desc = None
         out_desc = None
 
+    ACC_DTYPE = tl.int32 if A.dtype == torch.int8 else tl.float32
+
     _scaled_mm_per_token[grid](
         A=A,
         A_desc=A_desc,
@@ -305,12 +309,17 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]):
         B_IS_TRANS=B_is_trans,
         NEED_N_MASK=NEED_N_MASK,
         NEED_K_MASK=NEED_K_MASK,
+        ACC_DTYPE=ACC_DTYPE,
         **run_config,
     )
 
     return out
 
 
+fp8_scaled_mm_per_token = scaled_mm_per_token
+int8_scaled_mm_per_token = scaled_mm_per_token
+
+
 if __name__ == "__main__":
     import time
     import os
@@ -324,7 +333,7 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]):
     M_list = [1, 2, 4, 8, 16, 32, 48]
 
     print(f"{'='*80}")
-    print(f"Starting Autotune for FP8 Scaled MM (N={N}, K={K})")
+    print(f"Starting Autotune for Scaled MM (N={N}, K={K})")
     print(f"M values to test: {M_list}")
     print(f"Total configs per M: {len(get_test_configs())}")
     print(f"{'='*80}\n")
@@ -360,7 +369,7 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]):
     gt_C = d_A.mm(d_B)
 
     # 运行kernel验证正确性
-    fp8_scaled_mm_per_token(A_verify, B, Ascale_verify, Bscale, output_dtype, out_verify)
+    scaled_mm_per_token(A_verify, B, Ascale_verify, Bscale, output_dtype, out_verify)
 
     # 计算cosine similarity
     cosine_sim = F.cosine_similarity(out_verify.flatten().unsqueeze(0), gt_C.flatten().unsqueeze(0), dim=1)
@@ -390,7 +399,7 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]):
         A = test_data[M]["A"]
         Ascale = test_data[M]["Ascale"]
         out = test_data[M]["out"]
-        fp8_scaled_mm_per_token(A, B, Ascale, Bscale, output_dtype, out)
+        scaled_mm_per_token(A, B, Ascale, Bscale, output_dtype, out)
         print(f"[M={M}] Autotune completed!")
 
     Autotuner.end_autotune_warmup()
@@ -418,7 +427,7 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]):
         gt_C = d_A.mm(d_B)
 
         # 运行一次确保结果正确
-        fp8_scaled_mm_per_token(A, B, Ascale, Bscale, output_dtype, out)
+        scaled_mm_per_token(A, B, Ascale, Bscale, output_dtype, out)
         sgl_res = fp8_scaled_mm(A, B, Ascale, Bscale, output_dtype)
 
         cosine_sim = F.cosine_similarity(out.flatten().unsqueeze(0), gt_C.flatten().unsqueeze(0), dim=1)
@@ -437,7 +446,7 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]):
         ms_sgl = triton.testing.do_bench(fn_sgl, warmup=25, rep=100)
 
         # Our kernel
-        fn_ours = lambda: fp8_scaled_mm_per_token(A, B, Ascale, Bscale, output_dtype, out)
+        fn_ours = lambda: scaled_mm_per_token(A, B, Ascale, Bscale, output_dtype, out)
         ms_ours = triton.testing.do_bench_cudagraph(fn_ours, rep=100)
 
         print(f"[M={M}] BF16:       {ms_bf16:.3f} ms")
diff --git a/lightllm/common/basemodel/triton_kernel/quantize_gemm_int8.py b/lightllm/common/basemodel/triton_kernel/quantize_gemm_int8.py
deleted file mode 100644
index 4f3f6a385c..0000000000
--- a/lightllm/common/basemodel/triton_kernel/quantize_gemm_int8.py
+++ /dev/null
@@ -1,376 +0,0 @@
-import time
-import torch
-
-import triton
-import triton.language as tl
-
-
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_stages=2, num_warps=8),
-        triton.Config({}, num_stages=2, num_warps=4),
-        triton.Config({}, num_stages=2, num_warps=2),
-        triton.Config({}, num_stages=2, num_warps=1),
-     ],
-    key=['K'],
-)
-@triton.jit
-def quantize_int8_perrow_kernel(
-    fpa_ptr, a_ptr, as_ptr,
-    M, K, 
-    stride_fpam, stride_fpak,
-    stride_am, stride_ak,
-    stride_asm,
-    # Meta-parameters
-    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
-):
-    pid_m = tl.program_id(axis=0)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
-
-    fpa_ptrs = fpa_ptr + offs_am[:, None] * stride_fpam + offs_k[None, :] * stride_fpak
-    a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
-    a_max = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        fpa = tl.load(fpa_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
-        a_max = tl.maximum(a_max, tl.max(tl.abs(fpa), axis=1))
-        fpa_ptrs += BLOCK_SIZE_K * stride_fpak
-    a_scale = (a_max / 127.)
-    fpa_ptrs = fpa_ptr + offs_am[:, None] * stride_fpam + offs_k[None, :] * stride_fpak
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        fpa = tl.load(fpa_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
-        inta = (fpa / a_scale[:, None]).to(tl.int8)
-        tl.store(a_ptrs, inta, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K)
-        fpa_ptrs += BLOCK_SIZE_K * stride_fpak
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-    as_offs = pid_m * BLOCK_SIZE_M * stride_asm + tl.arange(0, BLOCK_SIZE_M)
-    tl.store(as_ptr + as_offs, a_scale)
-
-
-def quantize_int8_perrow(fpa):
-    a = torch.empty(fpa.shape, device=fpa.device, dtype=torch.int8)
-    a_scale = torch.empty(fpa.shape[0], device=fpa.device, dtype=fpa.dtype)
-    M, K = fpa.shape
-    BLOCK_SIZE_M = 1
-    BLOCK_SIZE_K = triton.next_power_of_2(K)
-    grid = (M // BLOCK_SIZE_M,)
-    quantize_int8_perrow_kernel[grid](
-        fpa, a, a_scale,
-        M, K,
-        fpa.stride(0), fpa.stride(1),
-        a.stride(0), a.stride(1),
-        a_scale.stride(0),
-        BLOCK_SIZE_M, BLOCK_SIZE_K,
-    )
-    return a, a_scale
-
-
-@triton.autotune(
-    configs=[
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 32,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),
-        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),
-	    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32,  'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-	    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-	    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
-	    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32,  'GROUP_SIZE_M': 16}, num_stages=4, num_warps=4),
-	    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 16}, num_stages=3, num_warps=8),
-	    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-	    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 16}, num_stages=4, num_warps=4),
-	    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 16}, num_stages=3, num_warps=8),
-	    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 32,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32,  'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32,  'GROUP_SIZE_M': 16}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 16}, num_stages=3, num_warps=8),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 16}, num_stages=4, num_warps=4),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 16}, num_stages=3, num_warps=8),
-		triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),
-    ],
-    key=['M', 'N', 'K'],
-    reset_to_zero=['c_ptr']
-)
-@triton.jit
-def matmul_kernel(
-    # Pointers to matrices
-    a_ptr, as_ptr, b_ptr, bs_ptr, c_ptr,
-    # Matrix dimensions
-    M, N, K,
-    # The stride variables represent how much to increase the ptr by when moving by 1
-    # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr`
-    # by to get the element one row down (A has M rows).
-    stride_am, stride_ak,
-    stride_asm,
-    stride_bk, stride_bn,
-    stride_bsn,
-    stride_cm, stride_cn,
-    # Meta-parameters
-    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
-    GROUP_SIZE_M: tl.constexpr, SPLIT_K: tl.constexpr, 
-):
-    """Kernel for computing the matmul C = A x B.
-    A has shape (M, K), B has shape (K, N) and C has shape (M, N)
-    """
-    # -----------------------------------------------------------
-    # Map program ids `pid` to the block of C it should compute.
-    # This is done in a grouped ordering to promote L2 data reuse.
-    # See above `L2 Cache Optimizations` section for details.
-    pid = tl.program_id(axis=0)
-    pid_sp_k = tl.program_id(axis=1)
-    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-    num_pid_in_group = GROUP_SIZE_M * num_pid_n
-    group_id = pid // num_pid_in_group
-    first_pid_m = group_id * GROUP_SIZE_M
-    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
-    pid_m = first_pid_m + (pid % group_size_m)
-    pid_n = (pid % num_pid_in_group) // group_size_m
-
-    # ----------------------------------------------------------
-    # Create pointers for the first blocks of A and B.
-    # We will advance this pointer as we move in the K direction
-    # and accumulate
-    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
-    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
-    # See above `Pointer Arithmetics` section for details
-    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
-    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
-    offs_k = pid_sp_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
-    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
-    as_ptrs = as_ptr + offs_am * stride_asm
-    bs_ptrs = bs_ptr + offs_bn * stride_bsn
-    a_scale = tl.load(as_ptrs, mask=offs_am < M, other=0.0)
-    b_scale = tl.load(bs_ptrs, mask=offs_bn < N, other=0.0)
-    # -----------------------------------------------------------
-    # Iterate to compute a block of the C matrix.
-    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
-    # of fp32 values for higher accuracy.
-    # `accumulator` will be converted back to fp16 after the loop.
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.int32)
-    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.
-        # If it is out of bounds, set it to 0.
-        a = tl.load(a_ptrs, mask=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 * SPLIT_K, other=0.0)
-        # We accumulate along the K dimension.
-        accumulator += tl.dot(a, b)
-        # Advance the ptrs to the next K block.
-        a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk
-    # You can fuse arbitrary activation functions here
-    # while the accumulator is still in FP32!
-    c = (accumulator.to(tl.float32) * a_scale[:, None] * b_scale[None, :]).to(c_ptr.dtype.element_ty)
-    # -----------------------------------------------------------
-    # Write back the block of the output matrix C with masks.
-    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-    if SPLIT_K == 1:
-        tl.store(c_ptrs, c, mask=c_mask)
-    else:
-        tl.atomic_add(c_ptrs, c, mask=c_mask)
-
-
-def matmul_quantize_int8(fpa, b, b_scale, out=None):
-    a, a_scale = quantize_int8_perrow(fpa)
-    # a, a_scale = quantize_int8(fpa, axis=1)
-    return matmul_int8(a, a_scale, b, b_scale, out)
-
-
-def matmul_int8(a, a_scale, b, b_scale, out=None):
-    # Check constraints.
-    assert a.shape[1] == b.shape[0], "Incompatible dimensions"
-    M, K = a.shape
-    K, N = b.shape
-    # Allocates output.
-    if out == None:
-        c = torch.zeros((M, N), device=a.device, dtype=torch.float16)
-    else:
-        c = out.fill_(0.)
-    grid = lambda META: (
-        triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),
-        META['SPLIT_K'],
-    )
-    matmul_kernel[grid](
-        a, a_scale, b, b_scale, c,
-        M, N, K,
-        a.stride(0), a.stride(1),
-        a_scale.stride(0),
-        b.stride(0), b.stride(1),
-        b_scale.stride(0),
-        c.stride(0), c.stride(1),
-    )
-    return c
-
-
-def quantize_int8(weight, axis=0, tp_rank=0):
-    # Weight shape: [H1, H2]
-    # Scale shape: [H2]
-    scale = weight.abs().amax(axis, keepdim=True) / 127.
-    weight = (weight / scale).to(torch.int8)
-    # col major will accelerate i8xi8 kernel.
-    if axis == 0:
-        weight = weight.t().contiguous().t()
-    scale = scale.squeeze(axis)
-    return weight.contiguous().cuda(tp_rank), scale.contiguous().cuda(tp_rank)
-
-
-def test_correct_int8(M=32, N=4096, K=4096):
-    a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-    b = torch.randn((K, N), device='cuda', dtype=torch.float16)
-    int_a, scale_a = quantize_int8_perrow(a)
-    cos = torch.nn.CosineSimilarity(0)
-    print("Quantization cos", cos((int_a * scale_a.unsqueeze(1)).flatten().to(torch.float32), a.flatten().to(torch.float32)))
-    int_b, scale_b = quantize_int8(b, axis=0)
-    triton_output = matmul_int8(int_a, scale_a, int_b, scale_b)
-    torch_output = torch.matmul(a, b)
-    print(f"triton_output={triton_output}")
-    print(f"torch_output={torch_output}")
-    cos = torch.nn.CosineSimilarity(0)
-    print("Output cos", cos(triton_output.flatten().to(torch.float32), torch_output.flatten().to(torch.float32)))
-
-
-def test_int8(M, K, N):
-    import time
-
-    print("M: {} K: {} N: {}".format(M, K, N))
-    torch.manual_seed(0)
-    a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-    b = torch.randn((K, N), device='cuda', dtype=torch.float16).contiguous()
-    int_b, scale_b = quantize_int8(b, axis=0)
-    for _ in range(10):
-        # int_a, a_scale = quantize_int8(a, 1)
-        int_a, a_scale = quantize_int8_perrow(a)
-        triton_output = matmul_int8(int_a, a_scale, int_b, scale_b)
-    torch.cuda.synchronize()
-    iters = 512
-    t1 = time.time()
-    for _ in range(iters):
-        #int_a, a_scale, _ = quantize_int8(a, 1)
-        int_a, a_scale = quantize_int8_perrow(a)
-    torch.cuda.synchronize()
-    qt2 = time.time()
-    for _ in range(iters):
-        triton_output = matmul_int8(int_a, a_scale, int_b, scale_b)
-    torch.cuda.synchronize()
-    t2 = time.time()
-    quant_time = qt2 - t1
-    triton_time = t2 - qt2
-    triton_tflops = 2 * M * N * K * 1e-12 / (triton_time / iters)
-    quant_bandwith = 2 * M * K * 1e-9 / (quant_time / iters)
-    print("Triton time cost: {} (tflops {}) + quant: {} (bandwidth {})".format(
-        triton_time, triton_tflops, quant_time, quant_bandwith))
-    for _ in range(10):
-        torch_output = torch.matmul(a, b)
-    torch.cuda.synchronize()
-    iters = 512
-    t1 = time.time()
-    for _ in range(iters):
-        torch_output = torch.matmul(a, b)
-    torch.cuda.synchronize()
-    t2 = time.time()
-    torch_time = t2 - t1
-    torch_tflops = 2 * M * N * K * 1e-12 / (torch_time / iters)
-    print("Torch time cost: {} (tflops {})".format(t2 - t1, torch_tflops))
-    return triton_time, torch_time, quant_time
-
-
-@triton.testing.perf_report(
-    triton.testing.Benchmark(
-        x_names=['M'],  # Argument names to use as an x-axis for the plot
-        x_vals=[32, 64, 128, 256] + [
-            512 * i * 2 for i in range(1, 17)
-        ],  # Different possible values for `x_name`
-        line_arg='provider',  # Argument name whose value corresponds to a different line in the plot
-        # Possible values for `line_arg`
-        line_vals=['cublas', 'triton-i8', 'triton-quant-i8', 'quant-perrow'],
-        # Label name for the lines
-        line_names=["cuBLAS", "Triton-i8", "Triton-Quant-i8", "Quant-perrow(GB/s)"],
-        # Line styles
-        styles=[('green', '-'), ('blue', '-'), ('red', '-'), ('purple', '-')],
-        ylabel="TFLOPS",  # Label name for the y-axis
-        plot_name="matmul-performance",  # Name for the plot, used also as a file name for saving the plot.
-        args={},
-    )
-)
-def benchmark(M, provider):
-    K = 10240
-    N = 27392 * 2 // 8
-    quantiles = [0.5, 0.2, 0.8]
-    if provider == 'cublas':
-        a = torch.randn((M, K), device='cuda', dtype=torch.float16)
-        b = torch.randn((K, N), device='cuda', dtype=torch.float16)
-        ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch.matmul(a, b), quantiles=quantiles)
-        perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
-    if provider == 'triton-i8':
-        a = torch.randn((M, K), device='cuda', dtype=torch.float16).to(torch.int8).contiguous()
-        b = torch.randn((K, N), device='cuda', dtype=torch.float16).to(torch.int8).contiguous()
-        int_a, a_scale = quantize_int8(a, axis=1)
-        int_b, b_scale = quantize_int8(b, axis=0)
-        ms, min_ms, max_ms = triton.testing.do_bench(lambda: matmul_int8(int_a, a_scale, int_b, b_scale), quantiles=quantiles)
-        perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
-    if provider == 'triton-quant-i8':
-        a = torch.randn((M, K), device='cuda', dtype=torch.float16).to(torch.int8).contiguous()
-        b = torch.randn((K, N), device='cuda', dtype=torch.float16).to(torch.int8).contiguous()
-        int_b, b_scale = quantize_int8(b, axis=0)
-        ms, min_ms, max_ms = triton.testing.do_bench(lambda: matmul_quantize_int8(a, int_b, b_scale), quantiles=quantiles)
-        perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
-    if provider == 'quant-perrow':
-        a = torch.randn((M, K), device='cuda', dtype=torch.float16).to(torch.int8).contiguous()
-        ms, min_ms, max_ms = triton.testing.do_bench(lambda: quantize_int8_perrow(a), quantiles=quantiles)
-        perf = lambda ms: 2 * M * K * 1e-9 / (ms * 1e-3)
-    return perf(ms), perf(min_ms), perf(max_ms)
-
-
-def test_model_layer(bs, sqe_len, hidden, inter, tp):
-    st1 = 0
-    st2 = 0
-    st3 = 0
-    t1, t2, t3 = test_int8(bs * sqe_len, hidden, hidden * 3 // tp)
-    st1 += t1
-    st2 += t2
-    st3 += t3
-    t1, t2, t3 = test_int8(bs * sqe_len, hidden // tp, hidden)
-    st1 += t1
-    st2 += t2
-    st3 += t3
-    t1, t2, t3 = test_int8(bs * sqe_len, hidden, inter * 2 // tp)
-    st1 += t1
-    st2 += t2
-    st3 += t3
-    t1, t2, t3 = test_int8(bs * sqe_len, inter // tp, hidden)
-    st1 += t1
-    st2 += t2
-    st3 += t3
-    print("Triton time {} Torch time {} Quant time {}".format(st1, st2, st3))
-
-
-if __name__ == "__main__":
-    test_correct_int8()
-    benchmark.run(show_plots=True, print_data=True)
-
-    bs = 32
-    hidden = 4096
-    inter  = 11008
-    prefill_len = 512
-    decode_len = 1
-    tp = 1
-    test_model_layer(bs, prefill_len, hidden, inter, tp)
-    test_model_layer(bs, decode_len, hidden, inter, tp)
diff --git a/lightllm/models/deepseek2/triton_kernel/repack_kv_index.py b/lightllm/common/basemodel/triton_kernel/repack_kv_index.py
similarity index 100%
rename from lightllm/models/deepseek2/triton_kernel/repack_kv_index.py
rename to lightllm/common/basemodel/triton_kernel/repack_kv_index.py
diff --git a/lightllm/common/kv_cache_mem_manager/__init__.py b/lightllm/common/kv_cache_mem_manager/__init__.py
index 66caf5d789..7d516e6728 100644
--- a/lightllm/common/kv_cache_mem_manager/__init__.py
+++ b/lightllm/common/kv_cache_mem_manager/__init__.py
@@ -1,20 +1,16 @@
 from .mem_manager import MemoryManager, ReadOnlyStaticsMemoryManager
-from .int8kv_mem_manager import INT8KVMemoryManager
 from .calibration_fp8kv_mem_manager import CalibrationFP8KVMemoryManager
 from .export_calibration_mem_manager import ExportCalibrationMemoryManager
 from .ppl_int8kv_mem_manager import PPLINT8KVMemoryManager
 from .ppl_int4kv_mem_manager import PPLINT4KVMemoryManager
 from .deepseek2_mem_manager import Deepseek2MemoryManager
-from .deepseek2_fp8kv_mem_manager import Deepseek2FP8KVMemoryManager
 
 __all__ = [
     "MemoryManager",
     "ReadOnlyStaticsMemoryManager",
-    "INT8KVMemoryManager",
     "CalibrationFP8KVMemoryManager",
     "ExportCalibrationMemoryManager",
     "PPLINT4KVMemoryManager",
     "PPLINT8KVMemoryManager",
     "Deepseek2MemoryManager",
-    "Deepseek2FP8KVMemoryManager",
 ]
diff --git a/lightllm/common/kv_cache_mem_manager/deepseek2_fp8kv_mem_manager.py b/lightllm/common/kv_cache_mem_manager/deepseek2_fp8kv_mem_manager.py
deleted file mode 100644
index 00699f4b15..0000000000
--- a/lightllm/common/kv_cache_mem_manager/deepseek2_fp8kv_mem_manager.py
+++ /dev/null
@@ -1,8 +0,0 @@
-import torch
-from .deepseek2_mem_manager import Deepseek2MemoryManager
-
-
-class Deepseek2FP8KVMemoryManager(Deepseek2MemoryManager):
-    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9):
-        # scale被追加到kv_buffer末尾, 因此加2, dtype统一改成uint8
-        super().__init__(size, torch.uint8, head_num, head_dim + 2, layer_num, always_copy, mem_fraction)
diff --git a/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py b/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py
index 7711734601..3d93e1b070 100644
--- a/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py
+++ b/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py
@@ -3,13 +3,14 @@
 import torch.distributed as dist
 from lightllm.server.pd_io_struct import KVMoveTask
 from .mem_manager import MemoryManager
-from typing import List, Union
+from typing import List, Union, Any
 from lightllm.utils.log_utils import init_logger
 from lightllm.common.kv_trans_kernel.kv_trans import kv_trans
 from lightllm.common.kv_trans_kernel.kv_trans_v2 import kv_trans_v2_for_d_node, kv_trans_v2_for_p_node
 from lightllm.distributed.pynccl import PyNcclCommunicator
 from lightllm.common.kv_trans_kernel.nixl_kv_trans import mla_page_io
 
+
 logger = init_logger(__name__)
 
 
@@ -17,6 +18,29 @@ class Deepseek2MemoryManager(MemoryManager):
     def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9):
         super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy, mem_fraction)
 
+    def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor):
+        """
+        将每一层生成的kv拷贝到mem manager对应mem_index 位置中
+        """
+        from ..basemodel.triton_kernel.kv_copy.mla_copy_kv import destindex_copy_kv
+
+        rope_dim = 64
+        kv_lora_rank = kv.shape[2] - rope_dim
+        assert kv_lora_rank + rope_dim == self.kv_buffer.shape[-1]
+
+        destindex_copy_kv(
+            kv[:, :, :kv_lora_rank],
+            kv[:, :, kv_lora_rank:],
+            mem_index,
+            self.kv_buffer[layer_index][:, :, :kv_lora_rank],
+            self.kv_buffer[layer_index][:, :, kv_lora_rank:],
+        )
+        return
+
+    def get_att_input_params(self, layer_index: int) -> Any:
+        kv = self.kv_buffer[layer_index]
+        return kv
+
     def get_cell_size(self):
         return self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype)
 
diff --git a/lightllm/common/kv_cache_mem_manager/export_calibration_mem_manager.py b/lightllm/common/kv_cache_mem_manager/export_calibration_mem_manager.py
index b2749176ea..ffdc9b2c94 100755
--- a/lightllm/common/kv_cache_mem_manager/export_calibration_mem_manager.py
+++ b/lightllm/common/kv_cache_mem_manager/export_calibration_mem_manager.py
@@ -1,6 +1,28 @@
+import torch
+from typing import Tuple, Any
 from .offline_fp8_quant_mem_manager import OfflineFP8QuantMemManager
 
 
 class ExportCalibrationMemoryManager(OfflineFP8QuantMemManager):
     def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9):
         super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy, mem_fraction, is_export_mode=True)
+
+    def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor):
+        """
+        将每一层生成的kv拷贝到mem manager对应mem_index 位置中
+        """
+        from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8
+
+        scales = self.scales
+        destindex_copy_kv_fp8(
+            kv,
+            mem_index,
+            scales[layer_index] if scales is not None else None,
+            self.kv_buffer[layer_index].view(torch.float8_e4m3fn),
+        )
+        return
+
+    def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]:
+        k = self.kv_buffer[layer_index][:, : self.head_num, :]
+        v = self.kv_buffer[layer_index][:, self.head_num :, :]
+        return k, v
diff --git a/lightllm/common/kv_cache_mem_manager/int8kv_mem_manager.py b/lightllm/common/kv_cache_mem_manager/int8kv_mem_manager.py
deleted file mode 100755
index 5725cdb7bb..0000000000
--- a/lightllm/common/kv_cache_mem_manager/int8kv_mem_manager.py
+++ /dev/null
@@ -1,29 +0,0 @@
-import torch
-
-from .mem_manager import MemoryManager
-
-
-class INT8KVMemoryManager(MemoryManager):
-    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True, mem_fraction=0.9):
-        self.kv_dtype = torch.int8
-        super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=True, mem_fraction=mem_fraction)
-
-    def get_cell_size(self):
-        return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(
-            self.kv_dtype
-        ) + 2 * self.head_num * self.layer_num * torch._utils._element_size(self.dtype)
-
-    def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
-        self.kv_buffer = torch.empty((layer_num, size + 1, 2 * head_num, head_dim), dtype=torch.int8, device="cuda")
-        self.scale_buffer = torch.empty((layer_num, size + 1, 2 * head_num, 1), dtype=dtype, device="cuda")
-
-    def _free_buffers(self):
-        self.kv_buffer = None
-        self.scale_buffer = None
-
-    def get_index_kv_buffer(self, index):
-        return {"kv_buffer": self.kv_buffer[:, index], "scale_buffer": self.scale_buffer[:, index]}
-
-    def load_index_kv_buffer(self, index, load_tensor_dict):
-        self.kv_buffer[:, index].copy_(load_tensor_dict["kv_buffer"])
-        self.scale_buffer[:, index].copy_(load_tensor_dict["scale_buffer"])
diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py
index d8fd93009f..1203cbdec7 100755
--- a/lightllm/common/kv_cache_mem_manager/mem_manager.py
+++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py
@@ -3,7 +3,7 @@
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
-from typing import List, Union
+from typing import List, Union, Tuple, Any
 from lightllm.common.kv_trans_kernel.kv_trans_v2 import kv_trans_for_dp
 from lightllm.server.pd_io_struct import KVMoveTask
 from lightllm.utils.log_utils import init_logger
@@ -21,6 +21,7 @@
 from multiprocessing.reduction import ForkingPickler
 from filelock import FileLock
 
+
 logger = init_logger(__name__)
 
 
@@ -65,6 +66,20 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False
         )
         self.HOLD_TOKEN_MEMINDEX = self.size
 
+    def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor):
+        """
+        将每一层生成的kv拷贝到mem manager对应mem_index 位置中
+        """
+        from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv
+
+        destindex_copy_kv(kv, mem_index, self.kv_buffer[layer_index])
+        return
+
+    def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]:
+        k = self.kv_buffer[layer_index][:, : self.head_num, :]
+        v = self.kv_buffer[layer_index][:, self.head_num :, :]
+        return k, v
+
     def get_cell_size(self):
         return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype)
 
diff --git a/lightllm/common/kv_cache_mem_manager/mem_utils.py b/lightllm/common/kv_cache_mem_manager/mem_utils.py
index 259c5a56f8..1ff58b89a0 100644
--- a/lightllm/common/kv_cache_mem_manager/mem_utils.py
+++ b/lightllm/common/kv_cache_mem_manager/mem_utils.py
@@ -1,12 +1,10 @@
 from . import (
     MemoryManager,
-    INT8KVMemoryManager,
     CalibrationFP8KVMemoryManager,
     ExportCalibrationMemoryManager,
     PPLINT8KVMemoryManager,
     PPLINT4KVMemoryManager,
     Deepseek2MemoryManager,
-    Deepseek2FP8KVMemoryManager,
 )
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.envs_utils import get_env_start_args
@@ -18,8 +16,6 @@
 
 @lru_cache(maxsize=None)
 def select_mem_manager_class():
-    mode = get_env_start_args().mode
-
     # case 1
     # 先判断是否是 deepseek 系列的模型
     model_class = get_llm_model_class()
@@ -27,38 +23,25 @@ def select_mem_manager_class():
 
     if issubclass(model_class, Deepseek2TpPartModel):
         mem_class = Deepseek2MemoryManager
-        if "triton_fp8kv" in mode:
-            mem_class = Deepseek2FP8KVMemoryManager
-
-        logger.info(f"Model kv cache using mode {mode}, mem_manager class: {mem_class}")
+        logger.info(f"Model kv cache using default, mem_manager class: {mem_class}")
         return mem_class
 
     # case normal
-    logger.info(f"mode setting params: {mode}")
-    if "ppl_int8kv" in mode or "ppl_int8kv_flashdecoding" in mode or "ppl_int8kv_flashdecoding_diverse" in mode:
+    logger.info(f"mode setting params: {get_env_start_args().llm_kv_type}")
+    if get_env_start_args().llm_kv_type == "int8kv":
         memory_manager_class = PPLINT8KVMemoryManager
-        logger.info(f"Model kv cache using mode {mode}")
-    elif "ppl_int4kv_flashdecoding" in mode:
+    elif get_env_start_args().llm_kv_type == "int4kv":
         memory_manager_class = PPLINT4KVMemoryManager
-        logger.info(f"Model kv cache using mode {mode}")
-    elif "triton_int8kv" in mode:
-        memory_manager_class = INT8KVMemoryManager
-        logger.info("Model kv cache using mode triton int8kv")
-    elif "triton_fp8kv" in mode:
-        raise Exception("currently only for deepseek")
-    elif "offline_calibration_fp8kv" in mode:
-        memory_manager_class = CalibrationFP8KVMemoryManager
-        logger.info("Model kv cache using mode offline calibration fp8kv")
-    elif "export_fp8kv_calibration" in mode:
+    elif get_env_start_args().llm_kv_type == "fp8kv":
         memory_manager_class = ExportCalibrationMemoryManager
-        logger.info("Using mode export fp8kv calibration")
-    else:
+    elif get_env_start_args().llm_kv_type == "None":
         memory_manager_class = MemoryManager
-        logger.info("Model kv cache using mode normal")
+
+    logger.info(f"Model kv cache using mem_manager class: {memory_manager_class}")
     return memory_manager_class
 
 
 @lru_cache(maxsize=None)
 def used_mem_manager_has_scale() -> bool:
     mem_class = select_mem_manager_class()
-    return mem_class in [PPLINT8KVMemoryManager, PPLINT4KVMemoryManager, INT8KVMemoryManager]
+    return mem_class in [PPLINT8KVMemoryManager, PPLINT4KVMemoryManager]
diff --git a/lightllm/common/kv_cache_mem_manager/offline_fp8_quant_mem_manager.py b/lightllm/common/kv_cache_mem_manager/offline_fp8_quant_mem_manager.py
index 5cc0b12d03..56a79a3b57 100755
--- a/lightllm/common/kv_cache_mem_manager/offline_fp8_quant_mem_manager.py
+++ b/lightllm/common/kv_cache_mem_manager/offline_fp8_quant_mem_manager.py
@@ -31,8 +31,10 @@ def __init__(
         self.scales_list = None
         self.abs_max = None
 
+        enable_fa3 = "fa3" in get_env_start_args().llm_prefill_att_backend
+
         if is_export_mode:
-            scales_shape = [layer_num, 2 * head_num] if get_env_start_args().enable_fa3 else [layer_num, 2]
+            scales_shape = [layer_num, 2 * head_num] if enable_fa3 else [layer_num, 2]
             self.abs_max = torch.zeros(scales_shape, dtype=torch.float32, device="cuda")
         elif get_env_start_args().kv_quant_calibration_config_path is not None:
             logger.info(
@@ -43,7 +45,7 @@ def __init__(
 
             self.scales_list = cfg["scales"]
             self.scales = torch.tensor(self.scales_list, dtype=torch.float32, device="cuda").view(cfg["scales_shape"])
-            if not get_env_start_args().enable_fa3:
+            if not enable_fa3:
                 self.scales = torch.repeat_interleave(self.scales, head_num, dim=-1)
             elif cfg["num_head"] > self.total_head_num:
                 factor = cfg["num_head"] // self.total_head_num
@@ -51,7 +53,7 @@ def __init__(
             elif cfg["num_head"] < self.total_head_num:
                 factor = self.total_head_num // cfg["num_head"]
                 self.scales = torch.repeat_interleave(self.scales, factor, dim=-1).contiguous()
-            if get_env_start_args().enable_fa3 and dist.is_initialized() and dist.get_world_size() > 1:
+            if enable_fa3 and dist.is_initialized() and dist.get_world_size() > 1:
                 half_head = self.total_head_num // 2
                 start_head = dist.get_rank() * head_num
                 end_head = start_head + head_num
@@ -65,6 +67,8 @@ def __init__(
             logger.warning("scales is None, no kv_quant_calibration_config_path be set, will use 1.0 as scales")
 
     def _load_and_check_config(self):
+        enable_fa3 = "fa3" in get_env_start_args().llm_prefill_att_backend
+
         if os.path.exists(get_env_start_args().kv_quant_calibration_config_path):
             with open(get_env_start_args().kv_quant_calibration_config_path, "r") as f:
                 cfg = json.load(f)
@@ -86,7 +90,7 @@ def _load_and_check_config(self):
                 raise ValueError(
                     f"num_head {cfg['num_head']} in config " f"not match current model head num {self.total_head_num}"
                 )
-            if get_env_start_args().enable_fa3:
+            if enable_fa3:
                 if cfg["quant_type"] != "per_head":
                     raise ValueError(f"quant type {cfg['num_head']} in config not match fa3 backend")
             else:
@@ -100,6 +104,7 @@ def _load_and_check_config(self):
             )
 
     def update_calibration_data(self, kv_buffer: torch.Tensor, layer_index: int):
+        enable_fa3 = "fa3" in get_env_start_args().llm_prefill_att_backend
         inference_counts = get_kv_quant_calibration_inference_count()
         warmup_counts = get_kv_quant_calibration_warmup_count()
         if not get_model_init_status() or self.count >= warmup_counts + inference_counts:
@@ -109,7 +114,7 @@ def update_calibration_data(self, kv_buffer: torch.Tensor, layer_index: int):
             logger.info("kv cache calibration mode will collect kv cache data for quantization calibration")
 
         if self.abs_max is not None and self.count >= warmup_counts:
-            if get_env_start_args().enable_fa3:
+            if enable_fa3:
                 kv_max = kv_buffer.abs().amax(dim=(0, 2)).to(torch.float32)
             else:
                 k_max = kv_buffer[:, : self.head_num, :].abs().amax(dim=()).to(torch.float32)
@@ -119,7 +124,7 @@ def update_calibration_data(self, kv_buffer: torch.Tensor, layer_index: int):
             if self.count == warmup_counts + inference_counts - 1 and layer_index == self.layer_num - 1:
                 final_abs_max = self.abs_max
                 if dist.is_initialized() and dist.get_world_size() > 1:
-                    if get_env_start_args().enable_fa3:
+                    if enable_fa3:
                         k_max, v_max = torch.chunk(self.abs_max, 2, dim=-1)
                         k_max = k_max.contiguous()
                         v_max = v_max.contiguous()
@@ -144,11 +149,13 @@ def update_calibration_data(self, kv_buffer: torch.Tensor, layer_index: int):
             self.count += 1
 
     def _export_calibration_data(self):
+        enable_fa3 = "fa3" in get_env_start_args().llm_prefill_att_backend
+
         model_arch = get_model_architectures(get_env_start_args().model_dir)
         cfg = {
             "version": "1.0",
             "architectures": model_arch,
-            "quant_type": "per_head" if get_env_start_args().enable_fa3 else "per_tensor",
+            "quant_type": "per_head" if enable_fa3 else "per_tensor",
             "qmin": self.qmin,
             "qmax": self.qmax,
             "num_layers": self.layer_num,
diff --git a/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py b/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py
index f3218594d3..559980dc12 100755
--- a/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py
+++ b/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py
@@ -1,5 +1,5 @@
 import torch
-
+from typing import Tuple, Any
 from .mem_manager import MemoryManager
 
 
@@ -9,6 +9,28 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True,
         self.group_quant_size = 8
         super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=always_copy, mem_fraction=mem_fraction)
 
+    def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor):
+        """
+        将每一层生成的kv拷贝到mem manager对应mem_index 位置中
+        """
+        from ..basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import destindex_copy_int4kv
+
+        destindex_copy_int4kv(
+            kv,
+            mem_index,
+            self.kv_buffer[layer_index],
+            self.scale_buffer[layer_index],
+            quant_group_size=self.group_quant_size,
+        )
+        return
+
+    def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]:
+        k = self.kv_buffer[layer_index][:, : self.head_num, :]
+        k_scale = self.scale_buffer[layer_index][:, : self.head_num, :]
+        v = self.kv_buffer[layer_index][:, self.head_num :, :]
+        v_scale = self.scale_buffer[layer_index][:, self.head_num :, :]
+        return (k, k_scale), (v, v_scale)
+
     def get_cell_size(self):
         return 2 * self.head_num * self.head_dim // 2 * self.layer_num * torch._utils._element_size(
             self.kv_dtype
diff --git a/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py b/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py
index 2a5aad7c8b..951d72e2c8 100755
--- a/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py
+++ b/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py
@@ -1,5 +1,5 @@
 import torch
-
+from typing import Tuple, Any
 from .mem_manager import MemoryManager
 
 
@@ -9,6 +9,28 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True,
         self.group_quant_size = 8
         super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=always_copy, mem_fraction=mem_fraction)
 
+    def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor):
+        """
+        将每一层生成的kv拷贝到mem manager对应mem_index 位置中
+        """
+        from ..basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import destindex_copy_quantize_kv
+
+        destindex_copy_quantize_kv(
+            kv,
+            mem_index,
+            self.kv_buffer[layer_index],
+            self.scale_buffer[layer_index],
+            quant_group_dim=self.group_quant_size,
+        )
+        return
+
+    def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]:
+        k = self.kv_buffer[layer_index][:, : self.head_num, :]
+        k_scale = self.scale_buffer[layer_index][:, : self.head_num, :]
+        v = self.kv_buffer[layer_index][:, self.head_num :, :]
+        v_scale = self.scale_buffer[layer_index][:, self.head_num :, :]
+        return (k, k_scale), (v, v_scale)
+
     def get_cell_size(self):
         return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(
             self.kv_dtype
diff --git a/lightllm/common/quantization/__init__.py b/lightllm/common/quantization/__init__.py
index 26f59258cd..1e47454490 100644
--- a/lightllm/common/quantization/__init__.py
+++ b/lightllm/common/quantization/__init__.py
@@ -1,11 +1,10 @@
 import yaml
 import collections
 from .registry import QUANTMETHODS
-from .torchao_quant import *
-from .w8a8_quant import *
-from .triton_quant.triton_quant import *
-from .deepgemm_quant import *
-from .awq_quant import *
+from .w8a8 import *
+from .deepgemm import *
+from .awq import *
+from .no_quant import *
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
@@ -37,7 +36,7 @@ def _mapping_quant_method(self):
         if self.hf_quantization_method == "fp8":
             block_size = self.hf_quantization_config.get("weight_block_size", None)
             if block_size == [128, 128]:
-                from lightllm.common.quantization.deepgemm_quant import HAS_DEEPGEMM
+                from lightllm.common.quantization.deepgemm import HAS_DEEPGEMM
 
                 if HAS_DEEPGEMM:
                     self.quant_type = "deepgemm-fp8w8a8-b128"
@@ -78,4 +77,6 @@ def get_quant_type(self, layer_num, name):
 
     def get_quant_method(self, layer_num, name):
         quant_type = self.get_quant_type(layer_num, name)
-        return QUANTMETHODS.get(quant_type)
+        quant_method = QUANTMETHODS.get(quant_type)
+        quant_method.hf_quantization_config = self.hf_quantization_config
+        return quant_method
diff --git a/lightllm/common/quantization/awq_quant.py b/lightllm/common/quantization/awq.py
similarity index 55%
rename from lightllm/common/quantization/awq_quant.py
rename to lightllm/common/quantization/awq.py
index 8c04cdcea9..ddb7674dd1 100644
--- a/lightllm/common/quantization/awq_quant.py
+++ b/lightllm/common/quantization/awq.py
@@ -1,34 +1,42 @@
-import os
 import torch
-from .quantize_method import QuantizationMethod
-from .registry import QUANTMETHODS
-import torch.nn.functional as F
-from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops, cutlass_scaled_mm
-from lightllm.utils.light_utils import HAS_LIGHTLLM_KERNEL, light_ops
-from typing import Any
-from typing import TYPE_CHECKING, Optional, Tuple
-from lightllm.utils.dist_utils import get_current_device_id
-
-if TYPE_CHECKING:
-    from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import MMWeightPack
+from typing import Any, Optional, Tuple
 
-if HAS_VLLM:
-    awq_dequantize = vllm_ops.awq_dequantize
-    awq_gemm = vllm_ops.awq_gemm
-    from vllm.model_executor.layers.quantization.utils.marlin_utils import (
-        check_marlin_supported,
-        marlin_permute_scales,
-        awq_to_marlin_zero_points,
-        should_use_atomic_add_reduce,
-        marlin_make_empty_g_idx,
-        marlin_make_workspace_new,
-    )
-    from vllm.scalar_type import scalar_types
-
-    TYPE_MAP = {
-        4: scalar_types.uint4,
-        8: scalar_types.uint8,
-    }
+from lightllm.common.quantization.quantize_method import QuantizationMethod, WeightPack
+from lightllm.common.quantization.registry import QUANTMETHODS
+from lightllm.utils.dist_utils import get_current_device_id
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+try:
+    from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops
+
+    if HAS_VLLM:
+        awq_dequantize = vllm_ops.awq_dequantize
+        awq_gemm = vllm_ops.awq_gemm
+        from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+            check_marlin_supported,
+            marlin_permute_scales,
+            awq_to_marlin_zero_points,
+            should_use_atomic_add_reduce,
+            marlin_make_empty_g_idx,
+            marlin_make_workspace_new,
+        )
+        from vllm.scalar_type import scalar_types
+
+        TYPE_MAP = {
+            4: scalar_types.uint4,
+            8: scalar_types.uint8,
+        }
+    else:
+        awq_dequantize = None
+        awq_gemm = None
+        TYPE_MAP = {}
+except ImportError:
+    HAS_VLLM = False
+    awq_dequantize = None
+    awq_gemm = None
+    TYPE_MAP = {}
 
 
 class AWQBaseQuantizationMethod(QuantizationMethod):
@@ -39,16 +47,17 @@ def __init__(self):
 
         self.cache_manager = g_cache_manager
 
-    def quantize(self, weight: torch.Tensor):
+    def quantize(self, weight: torch.Tensor, output: WeightPack, offset: int = 0):
         raise NotImplementedError("AWQ online quantization is not supported yet.")
 
     def apply(
         self,
         input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
+        weight_pack: WeightPack,
         out: Optional[torch.Tensor] = None,
         workspace: Optional[torch.Tensor] = None,
         use_custom_tensor_mananger: bool = True,
+        bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         raise NotImplementedError("AWQ online quantization is not supported yet.")
 
@@ -57,7 +66,7 @@ def method_name(self):
         return "awq-base"
 
 
-@QUANTMETHODS.register("awq")
+@QUANTMETHODS.register("awq", platform="cuda")
 class AWQW4A16QuantizationMethod(AWQBaseQuantizationMethod):
     def __init__(self):
         super().__init__()
@@ -72,21 +81,21 @@ def __init__(self):
     def method_name(self):
         return "awq"
 
-    def quantize(self, weight: torch.Tensor):
+    def quantize(self, weight: torch.Tensor, output: WeightPack, offset: int = 0):
         raise NotImplementedError("AWQ online quantization is not supported yet.")
 
     def apply(
         self,
         input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
+        weight_pack: WeightPack,
         out: Optional[torch.Tensor] = None,
         workspace: Optional[torch.Tensor] = None,
         use_custom_tensor_mananger: bool = True,
+        bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         qweight = weight_pack.weight
         weight_scale = weight_pack.weight_scale
         qzeros = weight_pack.weight_zero_point
-        bias = weight_pack.bias
 
         NEED_DEQUANT_WEIGHT = input_tensor.shape[:-1].numel() >= 256
         if NEED_DEQUANT_WEIGHT:
@@ -99,8 +108,35 @@ def apply(
             out.add_(bias)
         return out
 
-
-@QUANTMETHODS.register("awq_marlin")
+    def create_weight(
+        self, out_dim: int, in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
+    ) -> WeightPack:
+        group_size = self.hf_quantization_config["group_size"]
+        expert_prefix = (num_experts,) if num_experts > 1 else ()
+        weight = torch.empty(expert_prefix + (in_dim, out_dim // self.pack_factor), dtype=torch.int32).cuda(device_id)
+        weight_scale = torch.empty(expert_prefix + (in_dim // group_size, out_dim), dtype=dtype).cuda(device_id)
+        weight_zero_point = torch.empty(
+            expert_prefix + (in_dim // group_size, out_dim // self.pack_factor), dtype=torch.int32
+        ).cuda(device_id)
+        return WeightPack(weight=weight, weight_scale=weight_scale, weight_zero_point=weight_zero_point)
+
+    def load_weight(self, weight: torch.Tensor, weight_pack: WeightPack, start_idx: int) -> None:
+        start_idx = start_idx // self.pack_factor
+        weight_pack.weight[:, start_idx : start_idx + weight.shape[1]].copy_(weight)
+        return
+
+    def load_weight_scale(self, weight_scale: torch.Tensor, weight_pack: WeightPack, start_idx: int) -> None:
+        weight_pack.weight_scale[:, start_idx : start_idx + weight_scale.shape[1]].copy_(weight_scale)
+        return
+
+    def load_weight_zero_point(self, weight_zero_point: torch.Tensor, weight_pack: WeightPack, start_idx: int) -> None:
+        start_idx = start_idx // self.pack_factor
+        end_idx = start_idx + weight_zero_point.shape[1]
+        weight_pack.weight_zero_point[:, start_idx:end_idx].copy_(weight_zero_point)
+        return
+
+
+@QUANTMETHODS.register("awq_marlin", platform="cuda")
 class AWQMARLINW4A16QuantizationMethod(AWQBaseQuantizationMethod):
     def __init__(self):
         super().__init__()
@@ -115,20 +151,15 @@ def __init__(self):
         self.vllm_quant_type = TYPE_MAP[self.nbits]
         self.has_weight_scale = True
         self.has_weight_zero_point = True
+        self.tile_size = 16
 
     @property
     def method_name(self):
         return "awq_marlin"
 
-    def quantize(self, weight: torch.Tensor):
+    def quantize(self, weight: torch.Tensor, offset: int = 0) -> WeightPack:
         raise NotImplementedError("AWQ online quantization is not supported yet.")
 
-    def params_need_repack(self) -> bool:
-        """
-        用于说明是否需要对量化后的权重进行repack操作，目前只有awq支持
-        """
-        return True
-
     def params_repack(
         self, weight: torch.Tensor, weight_scale: torch.Tensor, weight_zero_point: torch.Tensor, dtype_type: torch.dtype
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -144,47 +175,18 @@ def params_repack(
         )
         return weight, weight_scale, weight_zero_point
 
-    def _process_weight_after_loading(self, weight: torch.Tensor) -> torch.Tensor:
-        assert self.hf_quantization_config is not None, "hf_quantization_config is not set"
-        self.k = weight.shape[0]
-        self.n = weight.shape[1] * self.pack_factor
-        return vllm_ops.awq_marlin_repack(
-            weight,
-            size_k=weight.shape[0],
-            size_n=weight.shape[1] * self.pack_factor,
-            num_bits=self.hf_quantization_config["bits"],
-        )
-
-    def _process_weight_scale_after_loading(self, weight_scale: torch.Tensor) -> torch.Tensor:
-        assert self.hf_quantization_config is not None, "hf_quantization_config is not set"
-        group_size = self.hf_quantization_config["group_size"]
-        return marlin_permute_scales(
-            weight_scale,
-            size_k=weight_scale.shape[0] * group_size,
-            size_n=weight_scale.shape[1],
-            group_size=self.hf_quantization_config["group_size"],
-        )
-
-    def _process_weight_zero_point_after_loading(self, weight_zero_point: torch.Tensor) -> torch.Tensor:
-        return awq_to_marlin_zero_points(
-            weight_zero_point,
-            size_k=weight_zero_point.shape[0],
-            size_n=weight_zero_point.shape[1] * self.pack_factor,
-            num_bits=self.hf_quantization_config["bits"],
-        )
-
     def apply(
         self,
         input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
+        weight_pack: WeightPack,
         out: Optional[torch.Tensor] = None,
         workspace: Optional[torch.Tensor] = None,
         use_custom_tensor_mananger: bool = True,
+        bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         qweight = weight_pack.weight
         weight_scale = weight_pack.weight_scale
         qzeros = weight_pack.weight_zero_point
-        bias = weight_pack.bias
         reshaped_x = input_tensor.reshape(-1, input_tensor.shape[-1])
 
         use_atomic_add = should_use_atomic_add_reduce(
@@ -219,6 +221,62 @@ def apply(
             out.add_(bias)
         return out
 
+    def create_weight(
+        self, out_dim: int, in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
+    ) -> WeightPack:
+        self.n = out_dim
+        self.k = in_dim
+        group_size = self.hf_quantization_config["group_size"]
+        expert_prefix = (num_experts,) if num_experts > 1 else ()
+        weight = torch.empty(
+            expert_prefix + (in_dim // self.tile_size, out_dim * self.tile_size // self.pack_factor), dtype=torch.int32
+        ).cuda(device_id)
+        weight_scale = torch.empty(expert_prefix + (in_dim // group_size, out_dim), dtype=dtype).cuda(device_id)
+        weight_zero_point = torch.empty(
+            expert_prefix + (in_dim // group_size, out_dim // self.pack_factor), dtype=torch.int32
+        ).cuda(device_id)
+        return WeightPack(weight=weight, weight_scale=weight_scale, weight_zero_point=weight_zero_point)
+
+    def load_weight(self, weight: torch.Tensor, weight_pack: WeightPack, start_idx: int) -> None:
+        assert self.hf_quantization_config is not None, "hf_quantization_config is not set"
+        device_id = get_current_device_id()
+        repack_weight = vllm_ops.awq_marlin_repack(
+            weight.cuda(device_id),
+            size_k=weight.shape[0],
+            size_n=weight.shape[1] * self.pack_factor,
+            num_bits=self.hf_quantization_config["bits"],
+        )
+        start_idx = start_idx // self.pack_factor * self.tile_size
+        weight_pack.weight[:, start_idx : start_idx + repack_weight.shape[1]].copy_(repack_weight)
+        return
+
+    def load_weight_scale(self, weight_scale: torch.Tensor, weight_pack: WeightPack, start_idx: int) -> None:
+        assert self.hf_quantization_config is not None, "hf_quantization_config is not set"
+        group_size = self.hf_quantization_config["group_size"]
+        device_id = get_current_device_id()
+        repack_weight_scale = marlin_permute_scales(
+            weight_scale.cuda(device_id),
+            size_k=weight_scale.shape[0] * group_size,
+            size_n=weight_scale.shape[1],
+            group_size=self.hf_quantization_config["group_size"],
+        )
+        weight_pack.weight_scale[:, start_idx : start_idx + repack_weight_scale.shape[1]].copy_(repack_weight_scale)
+        return
+
+    def load_weight_zero_point(self, weight_zero_point: torch.Tensor, weight_pack: WeightPack, start_idx: int) -> None:
+        device_id = get_current_device_id()
+        repack_weight_zero_point = awq_to_marlin_zero_points(
+            weight_zero_point.cuda(device_id),
+            size_k=weight_zero_point.shape[0],
+            size_n=weight_zero_point.shape[1] * self.pack_factor,
+            num_bits=self.hf_quantization_config["bits"],
+        )
+        start_idx = start_idx // self.pack_factor
+        weight_pack.weight_zero_point[:, start_idx : start_idx + repack_weight_zero_point.shape[1]].copy_(
+            repack_weight_zero_point
+        )
+        return
+
 
 # adapted from
 # https://github.com/vllm-project/vllm/blob/aef368aa08572505b820db01da82e2fbb3d43a72/vllm/model_executor/layers/quantization/awq_marlin.py#L211-L212
diff --git a/lightllm/common/quantization/deepgemm_quant.py b/lightllm/common/quantization/deepgemm.py
similarity index 51%
rename from lightllm/common/quantization/deepgemm_quant.py
rename to lightllm/common/quantization/deepgemm.py
index 7dbd3806b9..80be14c335 100644
--- a/lightllm/common/quantization/deepgemm_quant.py
+++ b/lightllm/common/quantization/deepgemm.py
@@ -1,20 +1,18 @@
-import os
 import torch
-from .quantize_method import QuantizationMethod
-from .registry import QUANTMETHODS
-import torch.nn.functional as F
-from lightllm.common.quantization.triton_quant.fp8.fp8act_quant_kernel import (
-    per_token_group_quant_fp8,
-    tma_align_input_scale,
-)
-from typing import TYPE_CHECKING, Optional
-
-if TYPE_CHECKING:
-    from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import MMWeightPack
+from typing import Optional
+
+from lightllm.common.quantization.quantize_method import QuantizationMethod, WeightPack
+from lightllm.common.quantization.registry import QUANTMETHODS
+from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
 try:
-    HAS_DEEPGEMM = True
     import deep_gemm
-except:
+
+    HAS_DEEPGEMM = True
+except ImportError:
     HAS_DEEPGEMM = False
 
 
@@ -26,17 +24,17 @@ def __init__(self):
         self.cache_manager = g_cache_manager
         assert HAS_DEEPGEMM, "deepgemm is not installed, you can't use quant api of it"
 
-    def quantize(self, weight: torch.Tensor):
-        """ """
-        pass
+    def quantize(self, weight: torch.Tensor, output: WeightPack, offset: int = 0):
+        raise NotImplementedError("Not implemented")
 
     def apply(
         self,
         input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
+        weight_pack: WeightPack,
         out: Optional[torch.Tensor] = None,
         workspace: Optional[torch.Tensor] = None,
         use_custom_tensor_mananger: bool = True,
+        bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         raise NotImplementedError("Not implemented")
 
@@ -45,12 +43,12 @@ def method_name(self):
         return "deepgemm-base"
 
 
-@QUANTMETHODS.register(["deepgemm-fp8w8a8-b128"])
+@QUANTMETHODS.register(["deepgemm-fp8w8a8-b128"], platform="cuda")
 class DeepGEMMFP8w8a8B128QuantizationMethod(DeepGEMMBaseQuantizationMethod):
     def __init__(self):
         super().__init__()
         self.block_size = 128
-        self.weight_suffix = None
+        self.weight_suffix = "weight"
         self.weight_zero_point_suffix = None
         self.weight_scale_suffix = "weight_scale_inv"
         self.has_weight_scale = True
@@ -60,26 +58,30 @@ def __init__(self):
     def method_name(self):
         return "deepgemm-fp8w8a8-b128"
 
-    def quantize(self, weight: torch.Tensor):
-        from lightllm.common.quantization.triton_quant.fp8.fp8w8a8_block_quant_kernel import weight_quant
+    def quantize(self, weight: torch.Tensor, output: WeightPack, offset: int = 0):
+        from lightllm.common.basemodel.triton_kernel.quantization.fp8w8a8_block_quant_kernel import weight_quant
 
-        weight, scale = weight_quant(weight, self.block_size)
-        return weight, scale, None
+        device = output.weight.device
+        weight, scale = weight_quant(weight.cuda(device), self.block_size)
+        output.weight[offset : offset + weight.shape[0], :].copy_(weight)
+        output.weight_scale[offset // self.block_size : offset + weight.shape[0] // self.block_size].copy_(scale)
+        return
 
     def apply(
         self,
         input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
+        weight_pack: "WeightPack",
         out: Optional[torch.Tensor] = None,
         workspace: Optional[torch.Tensor] = None,
         use_custom_tensor_mananger: bool = True,
+        bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         qweight = weight_pack.weight
         weight_scale = weight_pack.weight_scale
         input_scale = None
         alloc_func = torch.empty if not use_custom_tensor_mananger else self.cache_manager.empty
         m, k = input_tensor.shape
-        n = qweight.shape[1]
+        n = qweight.shape[0]
         if input_scale is None:
             qinput_tensor, input_scale = per_token_group_quant_fp8(
                 input_tensor,
@@ -92,9 +94,27 @@ def apply(
 
         if out is None:
             out = alloc_func((m, n), dtype=input_tensor.dtype, device=input_tensor.device)
-        _deepgemm_fp8_nt((qinput_tensor, input_scale), (qweight.t(), weight_scale.t()), out)
+        _deepgemm_fp8_nt((qinput_tensor, input_scale), (qweight, weight_scale), out)
         return out
 
+    def create_weight(
+        self, out_dim: int, in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
+    ) -> WeightPack:
+        expert_prefix = (num_experts,) if num_experts > 1 else ()
+        weight = torch.empty(expert_prefix + (out_dim, in_dim), dtype=torch.float8_e4m3fn).cuda(device_id)
+        scale_out_dim = (out_dim + self.block_size - 1) // self.block_size
+        scale_in_dim = (in_dim + self.block_size - 1) // self.block_size
+        weight_scale = torch.empty(expert_prefix + (scale_out_dim, scale_in_dim), dtype=torch.float32).cuda(device_id)
+        return WeightPack(weight=weight, weight_scale=weight_scale)
+
+    def load_weight(self, weight: torch.Tensor, weight_pack: WeightPack, start_idx: int) -> None:
+        weight_pack.weight[start_idx : start_idx + weight.shape[0]].copy_(weight)
+        return
+
+    def load_weight_scale(self, weight_scale: torch.Tensor, weight_pack: WeightPack, start_idx: int) -> None:
+        weight_pack.weight_scale[start_idx // self.block_size : start_idx + weight_scale.shape[0]].copy_(weight_scale)
+        return
+
 
 def _deepgemm_fp8_nt(a_tuple, b_tuple, out):
     if HAS_DEEPGEMM:
diff --git a/lightllm/common/quantization/no_quant.py b/lightllm/common/quantization/no_quant.py
new file mode 100644
index 0000000000..c05c90b210
--- /dev/null
+++ b/lightllm/common/quantization/no_quant.py
@@ -0,0 +1,56 @@
+import torch
+from typing import Optional
+
+from lightllm.common.quantization.quantize_method import QuantizationMethod, WeightPack
+from lightllm.common.quantization.registry import QUANTMETHODS
+
+
+@QUANTMETHODS.register("none", platform="musa")
+@QUANTMETHODS.register("none", platform="cuda")
+class NoQuantization(QuantizationMethod):
+    """No quantization - uses full precision weights."""
+
+    def apply(
+        self,
+        input_tensor: torch.Tensor,
+        weight_pack: WeightPack,
+        out: Optional[torch.Tensor] = None,
+        workspace: Optional[torch.Tensor] = None,
+        use_custom_tensor_mananger: bool = True,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
+
+        weight = weight_pack.weight.t()
+        if out is None:
+            shape = (input_tensor.shape[0], weight.shape[1])
+            dtype = input_tensor.dtype
+            device = input_tensor.device
+            if use_custom_tensor_mananger:
+                out = g_cache_manager.alloc_tensor(shape, dtype, device=device)
+            else:
+                out = torch.empty(shape, dtype=dtype, device=device)
+        if bias is None:
+            return torch.mm(input_tensor, weight, out=out)
+        return torch.addmm(bias, input_tensor, weight, out=out)
+
+    def create_weight(
+        self, out_dim: int, in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
+    ) -> WeightPack:
+        expert_prefix = (num_experts,) if num_experts > 1 else ()
+        weight = torch.empty(expert_prefix + (out_dim, in_dim), dtype=dtype).cuda(device_id)
+        return WeightPack(weight=weight, weight_scale=None, weight_zero_point=None)
+
+    def weight_need_quanted(self, weight: torch.Tensor) -> bool:
+        return False
+
+    def quantize(self, weight: torch.Tensor, output: WeightPack, offset: int = 0) -> None:
+        return
+
+    @property
+    def method_name(self):
+        return "none"
+
+    def load_weight(self, weight: torch.Tensor, weight_pack: WeightPack, start_idx: int = 0) -> None:
+        weight_pack.weight[start_idx : start_idx + weight.shape[0], :].copy_(weight)
+        return
diff --git a/lightllm/common/quantization/quantize_method.py b/lightllm/common/quantization/quantize_method.py
index 9b629bcaf1..4350307f1f 100644
--- a/lightllm/common/quantization/quantize_method.py
+++ b/lightllm/common/quantization/quantize_method.py
@@ -1,38 +1,62 @@
 import torch
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from lightllm.utils.dist_utils import get_current_device_id
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import Optional, Tuple
 
-if TYPE_CHECKING:
-    from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import MMWeightPack
+
+@dataclass
+class WeightPack:
+    weight: Optional[torch.Tensor] = None
+    weight_scale: Optional[torch.Tensor] = None
+    weight_zero_point: Optional[torch.Tensor] = None
+
+    def get_expert(self, expert_idx: int):
+        assert self.weight.ndim == 3, f"weight must be a 3D tensor, but got {self.weight.ndim}"
+        weight = self.weight[expert_idx]
+        weight_scale = self.weight_scale[expert_idx] if self.weight_scale is not None else None
+        weight_zero_point = self.weight_zero_point[expert_idx] if self.weight_zero_point is not None else None
+        return WeightPack(weight=weight, weight_scale=weight_scale, weight_zero_point=weight_zero_point)
+
+    def initialize_load_status(self, weight_num: int):
+        initial_loaded_status = [False, self.weight_scale is None, self.weight_zero_point is None]
+        self.load_ok = [initial_loaded_status.copy() for _ in range(weight_num)]
 
 
 class QuantizationMethod(ABC):
     def __init__(self):
         super().__init__()
         self.device_id_ = get_current_device_id()
-        self.weight_suffix = None
+        self.weight_suffix = "weight"
         self.weight_scale_suffix = None
         self.weight_zero_point_suffix = None
         self.act_scale_suffix = None
         self.has_weight_scale: bool = None
         self.has_weight_zero_point: bool = None
+        self.group_size: int = -1  # -1表示不分组即per-channel量化，其他表示分组大小
+        self.pack_factor: int = 1
+
         # 一些量化模式需要用到的额外量化参数，如awq量化
         self.hf_quantization_config = None
 
     @abstractmethod
-    def quantize(self, weights: torch.Tensor):
+    def quantize(
+        self,
+        weight: torch.Tensor,
+        output: WeightPack,
+        offset: int = 0,
+    ) -> None:
         pass
 
     @abstractmethod
     def apply(
         self,
         input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
-        bias: Optional[torch.Tensor] = None,
+        weight_pack: "WeightPack",
         out: Optional[torch.Tensor] = None,
         workspace: Optional[torch.Tensor] = None,
         use_custom_tensor_mananger: bool = True,
+        bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         pass
 
@@ -41,20 +65,26 @@ def apply(
     def method_name(self):
         pass
 
+    def create_weight(
+        self, out_dim: int, in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
+    ) -> WeightPack:
+        pass
+
     def weight_need_quanted(self, weight: torch.Tensor) -> bool:
         # 判断一个 weight 是否需要进行量化操作。
         return weight.dtype in [torch.bfloat16, torch.float16, torch.float32, torch.float64]
 
-    def params_need_repack(self) -> bool:
-        """
-        用于说明是否需要对量化后的权重进行repack操作，目前只有awq支持
-        """
-        return False
-
-    def params_repack(
-        self, weight: torch.Tensor, weight_scale: torch.Tensor, weight_zero_point: torch.Tensor, dtype_type: torch.dtype
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        一些量化方法在将参数完成量化后，为了加速性能，还需要将参数进行重拍，使算子性能达到最优，如awq方法。
-        """
-        return weight, weight_scale, weight_zero_point
+    def load_weight(self, weight: torch.Tensor, weight_pack: WeightPack, start_idx: int) -> None:
+        raise NotImplementedError(
+            f"quantization method {self.method_name} is not supported to load offline quantized weight"
+        )
+
+    def load_weight_scale(self, weight_scale: torch.Tensor, weight_pack: WeightPack, start_idx: int) -> None:
+        raise NotImplementedError(
+            f"quantization method {self.method_name} is not supported to load offline quantized weight scale"
+        )
+
+    def load_weight_zero_point(self, weight_zero_point: torch.Tensor, weight_pack: WeightPack, start_idx: int) -> None:
+        raise NotImplementedError(
+            f"quantization method {self.method_name} is not supported to load offline quantized weight zero point"
+        )
diff --git a/lightllm/common/quantization/registry.py b/lightllm/common/quantization/registry.py
index 674a22b60f..c9baa64e27 100644
--- a/lightllm/common/quantization/registry.py
+++ b/lightllm/common/quantization/registry.py
@@ -1,28 +1,31 @@
 from .quantize_method import QuantizationMethod
-from typing import Type
 
 
 class QuantMethodFactory:
     def __init__(self):
         self._quant_methods = {}
 
-    def register(self, names):
+    def register(self, names, platform="cuda"):
         def decorator(cls):
             local_names = names
             if isinstance(local_names, str):
                 local_names = [local_names]
             for n in local_names:
-                self._quant_methods[n] = cls
+                if n not in self._quant_methods:
+                    self._quant_methods[n] = {}
+                self._quant_methods[n][platform] = cls
             return cls
 
         return decorator
 
-    def get(self, key, *args, **kwargs) -> Type[QuantizationMethod]:
-        if key == "none":
-            return None
-        quant_method_class = self._quant_methods.get(key)
-        if not quant_method_class:
+    def get(self, key, platform="cuda", *args, **kwargs) -> "QuantizationMethod":
+        quant_method_class_dict = self._quant_methods.get(key)
+        if not quant_method_class_dict:
             raise ValueError(f"QuantMethod '{key}' not supported.")
+
+        quant_method_class = quant_method_class_dict.get(platform)
+        if quant_method_class is None:
+            raise ValueError(f"QuantMethod '{key}' for platform '{platform}' not supported.")
         return quant_method_class()
 
 
diff --git a/lightllm/common/quantization/torchao_quant.py b/lightllm/common/quantization/torchao_quant.py
deleted file mode 100644
index ba4115b1d9..0000000000
--- a/lightllm/common/quantization/torchao_quant.py
+++ /dev/null
@@ -1,169 +0,0 @@
-import os
-import torch
-from .quantize_method import QuantizationMethod
-from .registry import QUANTMETHODS
-import torch.nn.functional as F
-from typing import TYPE_CHECKING, Optional
-
-if TYPE_CHECKING:
-    from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import MMWeightPack
-
-try:
-    HAS_TORCH_AO = True
-    from torchao.quantization import (
-        int4_weight_only,
-        int8_weight_only,
-        float8_weight_only,
-        fpx_weight_only,
-        int8_dynamic_activation_int8_weight,
-        float8_dynamic_activation_float8_weight,
-        quantize_,
-    )
-    from torchao.utils import (
-        TORCH_VERSION_AT_LEAST_2_4,
-        TORCH_VERSION_AT_LEAST_2_5,
-    )
-except:
-    HAS_TORCH_AO = False
-
-
-class AOBaseQuantizationMethod(QuantizationMethod):
-    def __init__(self):
-        super().__init__()
-        assert HAS_TORCH_AO, "torchao is not installed, you can't use quant api of it"
-        assert TORCH_VERSION_AT_LEAST_2_4, "torchao requires torch >=2.4"
-        self.quant_func = None
-
-    def quantize(self, weight: torch.Tensor):
-        """ """
-        dummy_linear = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)
-        dummy_linear.weight = torch.nn.Parameter(weight.cuda(self.device_id_))
-        quantize_(dummy_linear, self.quant_func)
-        return dummy_linear.weight, None, None
-
-    def apply(
-        self,
-        input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
-        out: Optional[torch.Tensor] = None,
-        workspace: Optional[torch.Tensor] = None,
-        use_custom_tensor_mananger: bool = True,
-    ) -> torch.Tensor:
-        weights = weight_pack.weight
-        bias = weight_pack.bias
-        return F.linear(input_tensor, weights, bias)
-
-    @property
-    def method_name(self):
-        return "ao-base"
-
-
-@QUANTMETHODS.register(["ao-w4a16-256"])
-class AOW4A16QuantizationMethodGroup256(AOBaseQuantizationMethod):
-    def __init__(self):
-        super().__init__()
-        self.group_size = 256
-        self.quant_func = int4_weight_only(group_size=self.group_size)
-        self.has_weight_scale = False
-        self.has_weight_zero_point = False
-
-    @property
-    def method_name(self):
-        return "ao-w4a16-256"
-
-
-@QUANTMETHODS.register(["ao-w4a16-128"])
-class AOW4A16QuantizationMethodGroup128(AOBaseQuantizationMethod):
-    def __init__(self):
-        super().__init__()
-        self.group_size = 128
-        self.quant_func = int4_weight_only(group_size=self.group_size)
-        self.has_weight_scale = False
-        self.has_weight_zero_point = False
-
-    @property
-    def method_name(self):
-        return "ao-w4a16-128"
-
-
-@QUANTMETHODS.register(["ao-w4a16-64"])
-class AOW4A16QuantizationMethodGroup64(AOBaseQuantizationMethod):
-    def __init__(self):
-        super().__init__()
-        self.group_size = 64
-        self.quant_func = int4_weight_only(group_size=self.group_size)
-        self.has_weight_scale = False
-        self.has_weight_zero_point = False
-
-    @property
-    def method_name(self):
-        return "ao-w4a16-64"
-
-
-@QUANTMETHODS.register(["ao-w4a16-32"])
-class AOW4A16QuantizationMethodGroup32(AOBaseQuantizationMethod):
-    def __init__(self):
-        super().__init__()
-        self.group_size = 32
-        self.quant_func = int4_weight_only(group_size=self.group_size)
-        self.has_weight_scale = False
-        self.has_weight_zero_point = False
-
-    @property
-    def method_name(self):
-        return "ao-w4a16-32"
-
-
-@QUANTMETHODS.register("ao-w8a8")
-class AOW8A8QuantizationMethod(AOBaseQuantizationMethod):
-    def __init__(self):
-        super().__init__()
-        self.quant_func = int8_dynamic_activation_int8_weight()
-        self.has_weight_scale = False
-        self.has_weight_zero_point = False
-
-    @property
-    def method_name(self):
-        return "ao-w8a8"
-
-
-@QUANTMETHODS.register("ao-w8a16")
-class AOW8A16QuantizationMethod(AOBaseQuantizationMethod):
-    def __init__(self):
-        super().__init__()
-        self.quant_func = int8_weight_only()
-        self.has_weight_scale = False
-        self.has_weight_zero_point = False
-
-    @property
-    def method_name(self):
-        return "ao-w8a16"
-
-
-@QUANTMETHODS.register("ao-fp8w8a16")
-class AOFP8W8A16QuantizationMethod(AOBaseQuantizationMethod):
-    def __init__(self):
-        super().__init__()
-        is_cuda_8_9 = torch.cuda.is_available() and torch.cuda.get_device_capability() >= (8, 9)
-        assert is_cuda_8_9, "FP8 requires GPU with compute capability >= 8.9"
-        self.quant_func = float8_weight_only()
-        self.has_weight_scale = False
-        self.has_weight_zero_point = False
-
-    @property
-    def method_name(self):
-        return "ao-fp8w8a16"
-
-
-@QUANTMETHODS.register("ao-fp6w6a16")
-class AOFP6W6A16QuantizationMethod(AOBaseQuantizationMethod):
-    def __init__(self):
-        super().__init__()
-        assert TORCH_VERSION_AT_LEAST_2_5, "torchao fp6 requires torch >=2.5"
-        self.quant_func = fpx_weight_only(3, 2)
-        self.has_weight_scale = False
-        self.has_weight_zero_point = False
-
-    @property
-    def method_name(self):
-        return "ao-fp6w6a16"
diff --git a/lightllm/common/quantization/triton_quant/triton_quant.py b/lightllm/common/quantization/triton_quant/triton_quant.py
deleted file mode 100644
index 410f925a5e..0000000000
--- a/lightllm/common/quantization/triton_quant/triton_quant.py
+++ /dev/null
@@ -1,85 +0,0 @@
-import os
-import torch
-import torch.nn.functional as F
-from lightllm.common.quantization.quantize_method import QuantizationMethod
-from lightllm.common.quantization.registry import QUANTMETHODS
-from .fp8.fp8w8a8_block_gemm_kernel import w8a8_block_fp8_matmul
-from .fp8.fp8act_quant_kernel import per_token_group_quant_fp8
-from typing import TYPE_CHECKING, Optional
-
-if TYPE_CHECKING:
-    from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import MMWeightPack
-
-
-class TritonBaseQuantizationMethod(QuantizationMethod):
-    def __init__(self):
-        super().__init__()
-        from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
-
-        self.cache_manager = g_cache_manager
-
-    def quantize(self, weight: torch.Tensor):
-        pass
-
-    def apply(
-        self,
-        input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
-        out: Optional[torch.Tensor] = None,
-        workspace: Optional[torch.Tensor] = None,
-        use_custom_tensor_mananger: bool = True,
-    ) -> torch.Tensor:
-        raise NotImplementedError("Not implemented")
-
-
-@QUANTMETHODS.register(["triton-fp8w8a8-block128"])
-class TritonFP8w8a8QuantizationMethod(TritonBaseQuantizationMethod):
-    def __init__(self):
-        super().__init__()
-        self.is_moe = False
-        self.block_size = 128
-        self.weight_suffix = None
-        self.weight_zero_point_suffix = None
-        self.weight_scale_suffix = "weight_scale_inv"
-        self.has_weight_scale = True
-        self.has_weight_zero_point = False
-
-    def quantize(self, weight: torch.Tensor):
-        # TODO block-wise quant kernel
-        pass
-
-    def apply(
-        self,
-        input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
-        out: Optional[torch.Tensor] = None,
-        workspace: Optional[torch.Tensor] = None,
-        use_custom_tensor_mananger: bool = True,
-    ) -> torch.Tensor:
-        qweight = weight_pack.weight
-        weight_scale = weight_pack.weight_scale
-        input_scale = None
-        m, k = input_tensor.shape
-        n = qweight.shape[1]
-        alloc_func = torch.empty if not use_custom_tensor_mananger else self.cache_manager.empty
-        if input_scale is None:
-            input_tensor_q, input_scale = per_token_group_quant_fp8(
-                input_tensor, self.block_size, dtype=qweight.dtype, alloc_func=alloc_func
-            )
-        else:
-            # TODO
-            raise "statci input scale is not supported by triton fp8 block gemm kernel."
-        m = input_tensor.shape[0]
-        n = qweight.shape[1]
-        if out is None:
-            out = alloc_func((m, n), dtype=input_tensor.dtype, device=input_tensor.device)
-        w8a8_block_fp8_matmul(
-            input_tensor_q,
-            qweight,
-            input_scale,
-            weight_scale,
-            out,
-            (self.block_size, self.block_size),
-            dtype=input_tensor.dtype,
-        )
-        return out
diff --git a/lightllm/common/quantization/w8a8_quant.py b/lightllm/common/quantization/w8a8.py
similarity index 58%
rename from lightllm/common/quantization/w8a8_quant.py
rename to lightllm/common/quantization/w8a8.py
index 31004de4e3..0a74d9887c 100644
--- a/lightllm/common/quantization/w8a8_quant.py
+++ b/lightllm/common/quantization/w8a8.py
@@ -1,18 +1,17 @@
 import os
 import torch
-
-from lightllm.common.quantization.triton_quant.fp8.fp8w8a8_scaled_mm_per_token_kernel import fp8_scaled_mm_per_token
+import torch.nn.functional as F
+from typing import Optional
 from .quantize_method import QuantizationMethod
 from .registry import QUANTMETHODS
-import torch.nn.functional as F
-from typing import Optional, TYPE_CHECKING
-from lightllm.common.quantization.triton_quant.fp8.fp8act_quant_kernel import per_token_group_quant_fp8
-from lightllm.common.quantization.triton_quant.fp8.fp8w8a8_block_gemm_kernel import w8a8_block_fp8_matmul
+from lightllm.common.basemodel.triton_kernel.quantization.scaled_mm_per_token_kernel import fp8_scaled_mm_per_token
+from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8
+from lightllm.common.basemodel.triton_kernel.quantization.fp8w8a8_block_gemm_kernel import w8a8_block_fp8_matmul
 from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops, cutlass_scaled_mm
 from lightllm.utils.light_utils import HAS_LIGHTLLM_KERNEL, light_ops
 
-if TYPE_CHECKING:
-    from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import MMWeightPack
+
+from .quantize_method import WeightPack
 
 if HAS_LIGHTLLM_KERNEL:
 
@@ -38,16 +37,17 @@ def __init__(self):
 
         self.cache_manager = g_cache_manager
 
-    def quantize(self, weight: torch.Tensor):
-        pass
+    def quantize(self, weight: torch.Tensor, output: WeightPack, offset: int = 0) -> None:
+        raise NotImplementedError("Not implemented")
 
     def apply(
         self,
         input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
+        weight_pack: WeightPack,
         out: Optional[torch.Tensor] = None,
         workspace: Optional[torch.Tensor] = None,
         use_custom_tensor_mananger: bool = True,
+        bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         raise NotImplementedError("Not implemented")
 
@@ -55,35 +55,40 @@ def apply(
     def method_name(self):
         return "w8a8-base"
 
+    def create_weight(
+        self, out_dim: int, in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
+    ) -> WeightPack:
+        raise NotImplementedError("Not implemented")
+
 
-@QUANTMETHODS.register(["vllm-w8a8", "w8a8"])
+@QUANTMETHODS.register(["vllm-w8a8", "w8a8"], platform="cuda")
 class w8a8QuantizationMethod(BaseQuantizationMethod):
     def __init__(self):
         super().__init__()
         self.has_weight_scale = True
         self.has_weight_zero_point = False
 
-    def quantize(self, weight: torch.Tensor):
-        if isinstance(weight, tuple):
-            return (weight[0].transpose(0, 1).cuda(self.device_id_),) + weight[1:]
-        weight = weight.float()
+    def quantize(self, weight: torch.Tensor, output: WeightPack, offset: int = 0) -> None:
+        weight = weight.float().cuda(self.device_id_)
         scale = weight.abs().max(dim=-1)[0] / 127
-        weight = weight.transpose(0, 1) / scale.reshape(1, -1)
+        weight = weight / scale.reshape(-1, 1)
         weight = torch.round(weight.clamp(min=-128, max=127)).to(dtype=torch.int8)
-        return weight.cuda(self.device_id_), scale.cuda(self.device_id_), None
+        output.weight[offset : offset + weight.shape[0]].copy_(weight)
+        output.weight_scale[offset : offset + weight.shape[0]].copy_(scale)
+        return
 
     def apply(
         self,
         input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
+        weight_pack: WeightPack,
         out: Optional[torch.Tensor] = None,
         workspace: Optional[torch.Tensor] = None,
         use_custom_tensor_mananger: bool = True,
+        bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         input_scale = None
-        qweight = weight_pack.weight
+        qweight = weight_pack.weight.t()
         weight_scale = weight_pack.weight_scale
-        bias = weight_pack.bias
         input_scale = None  # dynamic quantization for input tensor
         x_q, x_scale, x_zp = vllm_ops.scaled_int8_quant(input_tensor, scale=input_scale, azp=None, symmetric=True)
         m = input_tensor.shape[0]
@@ -100,8 +105,16 @@ def apply(
     def method_name(self):
         return "vllm-w8a8"
 
+    def create_weight(
+        self, out_dim: int, in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
+    ) -> WeightPack:
+        expert_prefix = (num_experts,) if num_experts > 1 else ()
+        weight = torch.empty(expert_prefix + (out_dim, in_dim), dtype=torch.int8).cuda(device_id)
+        weight_scale = torch.empty(expert_prefix + (out_dim,), dtype=torch.float32).cuda(device_id)
+        return WeightPack(weight=weight, weight_scale=weight_scale)
 
-@QUANTMETHODS.register(["vllm-fp8w8a8", "fp8w8a8"])
+
+@QUANTMETHODS.register(["vllm-fp8w8a8", "fp8w8a8"], platform="cuda")
 class FP8w8a8QuantizationMethod(BaseQuantizationMethod):
     def __init__(self):
         super().__init__()
@@ -109,19 +122,20 @@ def __init__(self):
         self.has_weight_scale = True
         self.has_weight_zero_point = False
 
-    def quantize(self, weight: torch.Tensor):
+    def quantize(self, weight: torch.Tensor, output: WeightPack, offset: int = 0) -> None:
         if self.is_moe:
-            return self.quantize_moe(weight)
+            return self.quantize_moe(weight, output, offset)
         qweight, weight_scale = scaled_fp8_quant(
-            weight.contiguous().cuda(self.device_id_), scale=None, use_per_token_if_dynamic=True
+            weight.cuda(self.device_id_), scale=None, use_per_token_if_dynamic=True
         )
-        return qweight.transpose(0, 1), weight_scale, None
+        output.weight[offset : offset + qweight.shape[0], :].copy_(qweight)
+        output.weight_scale[offset : offset + weight_scale.shape[0]].copy_(weight_scale.view(-1))
+        return
 
-    def quantize_moe(self, weight: torch.Tensor):
+    def quantize_moe(self, weight: torch.Tensor) -> WeightPack:
         num_experts = weight.shape[0]
-        qweights = []
-        weight_scales = []
         qweights = torch.empty_like(weight, dtype=torch.float8_e4m3fn).cuda(self.device_id_)
+        weight_scales = []
         for i in range(num_experts):
             qweight, weight_scale = scaled_fp8_quant(
                 weight[i].contiguous().cuda(self.device_id_), scale=None, use_per_token_if_dynamic=True
@@ -129,19 +143,19 @@ def quantize_moe(self, weight: torch.Tensor):
             qweights[i] = qweight
             weight_scales.append(weight_scale)
         weight_scale = torch.stack(weight_scales, dim=0).contiguous()
-        return qweights, weight_scale, None
+        return WeightPack(weight=qweights, weight_scale=weight_scale)
 
     def apply(
         self,
         input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
+        weight_pack: WeightPack,
         out: Optional[torch.Tensor] = None,
         workspace: Optional[torch.Tensor] = None,
         use_custom_tensor_mananger: bool = True,
+        bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        qweight = weight_pack.weight
+        qweight = weight_pack.weight.t()
         weight_scale = weight_pack.weight_scale
-        bias = weight_pack.bias
         x_q, x_scale = scaled_fp8_quant(input_tensor, scale=None, scale_ub=None, use_per_token_if_dynamic=True)
         m = input_tensor.shape[0]
         n = qweight.shape[1]
@@ -160,8 +174,16 @@ def apply(
     def method_name(self):
         return "vllm-fp8w8a8"
 
+    def create_weight(
+        self, out_dim: int, in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
+    ) -> WeightPack:
+        expert_prefix = (num_experts,) if num_experts > 1 else ()
+        weight = torch.empty(expert_prefix + (out_dim, in_dim), dtype=torch.float8_e4m3fn).cuda(device_id)
+        weight_scale = torch.empty(expert_prefix + (out_dim,), dtype=torch.float32).cuda(device_id)
+        return WeightPack(weight=weight, weight_scale=weight_scale)
 
-@QUANTMETHODS.register(["vllm-fp8w8a8-b128", "fp8w8a8-b128"])
+
+@QUANTMETHODS.register(["vllm-fp8w8a8-b128", "fp8w8a8-b128"], platform="cuda")
 class FP8w8a8B128QuantizationMethod(BaseQuantizationMethod):
     def __init__(self):
         super().__init__()
@@ -170,21 +192,26 @@ def __init__(self):
         self.has_weight_scale = True
         self.has_weight_zero_point = False
 
-    def quantize(self, weight: torch.Tensor):
+    def quantize(self, weight: torch.Tensor, output: WeightPack, offset: int = 0) -> None:
+        from lightllm.common.quantization.triton_quant.fp8.fp8w8a8_block_quant_kernel import weight_quant
 
-        raise Exception("Not implemented")
+        device = output.weight.device
+        weight, scale = weight_quant(weight.cuda(device), self.block_size)
+        output.weight[offset : offset + weight.shape[0], :].copy_(weight)
+        output.weight_scale[offset // self.block_size : offset + weight.shape[0] // self.block_size].copy_(scale)
+        return
 
     def apply(
         self,
         input_tensor: torch.Tensor,
-        weight_pack: "MMWeightPack",
+        weight_pack: WeightPack,
         out: Optional[torch.Tensor] = None,
         workspace: Optional[torch.Tensor] = None,
         use_custom_tensor_mananger: bool = True,
+        bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        qweight = weight_pack.weight
-        weight_scale = weight_pack.weight_scale
-        bias = weight_pack.bias
+        qweight = weight_pack.weight.t()
+        weight_scale = weight_pack.weight_scale.t()
         input_scale = None  # dynamic quantization for input tensor
         m, k = input_tensor.shape
         n = qweight.shape[1]
@@ -213,3 +240,13 @@ def apply(
     @property
     def method_name(self):
         return "vllm-fp8w8a8-b128"
+
+    def create_weight(
+        self, out_dim: int, in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
+    ) -> WeightPack:
+        expert_prefix = (num_experts,) if num_experts > 1 else ()
+        weight = torch.empty(expert_prefix + (out_dim, in_dim), dtype=torch.float8_e4m3fn).cuda(device_id)
+        weight_scale = torch.empty(
+            expert_prefix + (out_dim // self.block_size, in_dim // self.block_size), dtype=torch.float32
+        ).cuda(device_id)
+        return WeightPack(weight=weight, weight_scale=weight_scale)
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
new file mode 100644
index 0000000000..9f44ee6c30
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_4090_D.json
@@ -0,0 +1,22 @@
+{
+  "16_8192": {
+    "BLOCK_N": 32,
+    "num_stages": 2,
+    "num_warps": 4
+  },
+  "32_8192": {
+    "BLOCK_N": 16,
+    "num_stages": 2,
+    "num_warps": 2
+  },
+  "64_8192": {
+    "BLOCK_N": 16,
+    "num_stages": 2,
+    "num_warps": 2
+  },
+  "8_8192": {
+    "BLOCK_N": 32,
+    "num_stages": 2,
+    "num_warps": 8
+  }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/fp8_scaled_mm_per_token:v3/{K=14336,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/scaled_mm_per_token:v1/{K=14336,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/fp8_scaled_mm_per_token:v3/{K=14336,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/scaled_mm_per_token:v1/{K=14336,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/fp8_scaled_mm_per_token:v3/{K=4096,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/scaled_mm_per_token:v1/{K=4096,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/fp8_scaled_mm_per_token:v3/{K=4096,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/scaled_mm_per_token:v1/{K=4096,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/fp8_scaled_mm_per_token:v3/{K=5120,N=2048,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/scaled_mm_per_token:v1/{K=5120,N=2048,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/fp8_scaled_mm_per_token:v3/{K=5120,N=2048,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/scaled_mm_per_token:v1/{K=5120,N=2048,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/fp8_scaled_mm_per_token:v3/{K=5120,N=28672,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/scaled_mm_per_token:v1/{K=5120,N=28672,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/fp8_scaled_mm_per_token:v3/{K=5120,N=28672,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/scaled_mm_per_token:v1/{K=5120,N=28672,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/fp8_scaled_mm_per_token:v3/{K=5120,N=4096,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/scaled_mm_per_token:v1/{K=5120,N=4096,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/fp8_scaled_mm_per_token:v3/{K=5120,N=4096,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_4090_D/scaled_mm_per_token:v1/{K=5120,N=4096,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_4090_D.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
new file mode 100644
index 0000000000..4fa2f949f2
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/_fwd_kernel_flash_decode_diverse_stage2:v1/{block_seq=256,gqa_group_size=4,out_dtype=torch.bfloat16,q_head_dim=128}_NVIDIA_GeForce_RTX_5090.json
@@ -0,0 +1,22 @@
+{
+  "16_8192": {
+    "BLOCK_N": 32,
+    "num_stages": 2,
+    "num_warps": 4
+  },
+  "32_8192": {
+    "BLOCK_N": 32,
+    "num_stages": 4,
+    "num_warps": 4
+  },
+  "64_8192": {
+    "BLOCK_N": 16,
+    "num_stages": 3,
+    "num_warps": 2
+  },
+  "8_8192": {
+    "BLOCK_N": 16,
+    "num_stages": 2,
+    "num_warps": 2
+  }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=13824,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=13824,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=13824,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=13824,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=14336,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=14336,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=14336,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=14336,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=1536,N=1536,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=1536,N=1536,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=1536,N=1536,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=1536,N=1536,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=1536,N=8960,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=1536,N=8960,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=1536,N=8960,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=1536,N=8960,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=4096,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=4096,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=4096,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=4096,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=5120,N=13824,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=5120,N=13824,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=5120,N=13824,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=5120,N=13824,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=5120,N=2048,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=5120,N=2048,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=5120,N=2048,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=5120,N=2048,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=5120,N=28672,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=5120,N=28672,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=5120,N=28672,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=5120,N=28672,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=5120,N=4096,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=5120,N=4096,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=5120,N=4096,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=5120,N=4096,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=5120,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=5120,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=5120,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=5120,N=5120,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=8960,N=1536,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=8960,N=1536,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{K=8960,N=1536,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{K=8960,N=1536,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{N=14336,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{N=14336,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
similarity index 100%
rename from lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/fp8_scaled_mm_per_token:v3/{N=14336,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
rename to lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_GeForce_RTX_5090/scaled_mm_per_token:v1/{N=14336,out_dtype=torch.bfloat16}_NVIDIA_GeForce_RTX_5090.json
diff --git a/lightllm/distributed/communication_op.py b/lightllm/distributed/communication_op.py
index d5c96f821a..d606d757c1 100644
--- a/lightllm/distributed/communication_op.py
+++ b/lightllm/distributed/communication_op.py
@@ -136,9 +136,9 @@ def get_group(self, group_index: int) -> CustomProcessGroup:
         return self.groups[group_index]
 
     def new_deepep_group(self, n_routed_experts, hidden_size):
-        moe_mode = os.getenv("MOE_MODE", "TP")
+        enable_ep_moe = get_env_start_args().enable_ep_moe
         num_max_dispatch_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank()
-        if moe_mode == "TP":
+        if not enable_ep_moe:
             self.ep_buffer = None
             return
         assert HAS_DEEPEP, "deep_ep is required for expert parallelism"
diff --git a/lightllm/models/__init__.py b/lightllm/models/__init__.py
index 4ee02f003b..539b32decb 100644
--- a/lightllm/models/__init__.py
+++ b/lightllm/models/__init__.py
@@ -1,4 +1,3 @@
-from lightllm.models.cohere.model import CohereTpPartModel
 from lightllm.models.mixtral.model import MixtralTpPartModel
 from lightllm.models.bloom.model import BloomTpPartModel
 from lightllm.models.llama.model import LlamaTpPartModel
@@ -8,7 +7,6 @@
 from lightllm.models.qwen2.model import Qwen2TpPartModel
 from lightllm.models.qwen3.model import Qwen3TpPartModel
 from lightllm.models.qwen3_moe.model import Qwen3MOEModel
-from lightllm.models.chatglm2.model import ChatGlm2TpPartModel
 from lightllm.models.internlm.model import InternlmTpPartModel
 from lightllm.models.stablelm.model import StablelmTpPartModel
 from lightllm.models.internlm2.model import Internlm2TpPartModel
diff --git a/lightllm/models/bloom/layer_infer/post_layer_infer.py b/lightllm/models/bloom/layer_infer/post_layer_infer.py
index 7938869f5a..ec1d94458f 100644
--- a/lightllm/models/bloom/layer_infer/post_layer_infer.py
+++ b/lightllm/models/bloom/layer_infer/post_layer_infer.py
@@ -10,12 +10,10 @@
 class BloomPostLayerInfer(LlamaPostLayerInfer):
     """ """
 
-    def __init__(self, network_config, mode):
+    def __init__(self, network_config):
         repair_config(config=network_config, same_names=["layer_norm_epsilon", "rms_norm_eps"])
-        super().__init__(network_config, mode)
+        super().__init__(network_config)
         return
 
     def _norm(self, input, infer_state, layer_weight: BloomPreAndPostLayerWeight) -> torch.Tensor:
-        return layer_weight.final_norm_weight_.layernorm_forward(
-            input=input, eps=self.eps_, alloc_func=self.alloc_tensor
-        )
+        return layer_weight.final_norm_weight_(input=input, eps=self.eps_, alloc_func=self.alloc_tensor)
diff --git a/lightllm/models/bloom/layer_infer/pre_layer_infer.py b/lightllm/models/bloom/layer_infer/pre_layer_infer.py
index baf1d3084d..e84069e116 100644
--- a/lightllm/models/bloom/layer_infer/pre_layer_infer.py
+++ b/lightllm/models/bloom/layer_infer/pre_layer_infer.py
@@ -9,23 +9,23 @@
 class BloomPreLayerInfer(PreLayerInferTpl):
     """ """
 
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         self.eps_ = network_config["layer_norm_epsilon"]
         return
 
     def _norm(self, input, infer_state, layer_weight: BloomPreAndPostLayerWeight) -> torch.Tensor:
-        return layer_weight.pre_norm_weight_.layernorm_forward(input=input, eps=self.eps_, alloc_func=self.alloc_tensor)
+        return layer_weight.pre_norm_weight_(input=input, eps=self.eps_, alloc_func=self.alloc_tensor)
 
     def context_forward(self, input_ids, infer_state: InferStateInfo, layer_weight: BloomPreAndPostLayerWeight):
-        input_embdings = layer_weight.wte_weight_.embedding(input_ids=input_ids, alloc_func=self.alloc_tensor)
+        input_embdings = layer_weight.wte_weight_(input_ids=input_ids, alloc_func=self.alloc_tensor)
         if self.tp_world_size_ > 1:
             all_reduce(input_embdings, group=infer_state.dist_group, op=dist.ReduceOp.SUM, async_op=False)
         input_embdings = self._norm(input_embdings, infer_state, layer_weight)
         return input_embdings
 
     def token_forward(self, input_ids, infer_state: InferStateInfo, layer_weight: BloomPreAndPostLayerWeight):
-        input_embdings = layer_weight.wte_weight_.embedding(input_ids=input_ids, alloc_func=self.alloc_tensor)
+        input_embdings = layer_weight.wte_weight_(input_ids=input_ids, alloc_func=self.alloc_tensor)
         if self.tp_world_size_ > 1:
             all_reduce(input_embdings, group=infer_state.dist_group, op=dist.ReduceOp.SUM, async_op=False)
         input_embdings = self._norm(input_embdings, infer_state, layer_weight)
diff --git a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py
index d82a23d039..60d584eebd 100755
--- a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py
@@ -2,16 +2,15 @@
 from typing import Tuple
 from lightllm.common.basemodel import TransformerLayerInferTpl
 from lightllm.models.bloom.layer_weights.transformer_layer_weight import BloomTransformerLayerWeight
-from lightllm.models.bloom.triton_kernel.context_flashattention_nopad import context_attention_fwd
-from lightllm.models.bloom.triton_kernel.token_flashattention_nopad import token_attention_fwd
 from lightllm.common.basemodel import InferStateInfo
+from lightllm.common.basemodel.attention.base_att import AttControl
 
 
 class BloomTransformerLayerInfer(TransformerLayerInferTpl):
     """ """
 
-    def __init__(self, layer_num, network_config, mode):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         self.eps_ = network_config["layer_norm_epsilon"]
         self.tp_q_head_num_ = network_config["num_attention_heads"] // self.tp_world_size_
         self.tp_k_head_num_ = self.tp_q_head_num_
@@ -21,17 +20,51 @@ def __init__(self, layer_num, network_config, mode):
         self.embed_dim_ = network_config["n_embed"]
         return
 
+    def _context_attention_kernel(
+        self,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        infer_state: InferStateInfo,
+        layer_weight: BloomTransformerLayerWeight,
+        out=None,
+    ) -> torch.Tensor:
+        _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_)
+        _q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
+        o_tensor = infer_state.prefill_att_state.prefill_att(
+            q=_q,
+            k=_k,
+            v=_v,
+            att_control=AttControl(use_alibi=True, tp_alibi=layer_weight.tp_alibi),
+            alloc_func=self.alloc_tensor,
+        )
+        o_tensor = o_tensor.view(q.shape)
+        return o_tensor
+
+    def _token_attention_kernel(
+        self, q: torch.Tensor, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight, out=None
+    ) -> torch.Tensor:
+        _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_)
+        _q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
+        o_tensor = infer_state.decode_att_state.decode_att(
+            q=_q,
+            k=_k,
+            v=_v,
+            att_control=AttControl(use_alibi=True, tp_alibi=layer_weight.tp_alibi),
+            alloc_func=self.alloc_tensor,
+        )
+        return o_tensor.view(q.shape)
+
     def _att_norm(
         self, input: torch.Tensor, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight
     ) -> torch.Tensor:
-        return layer_weight.att_norm_weight_.layernorm_forward(
+        return layer_weight.att_norm_weight_(
             input=input.view(-1, self.embed_dim_), eps=self.eps_, alloc_func=self.alloc_tensor
         )
 
     def _ffn_norm(
         self, input: torch.Tensor, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight
     ) -> torch.Tensor:
-        return layer_weight.ffn_norm_weight_.layernorm_forward(
+        return layer_weight.ffn_norm_weight_(
             input=input.view(-1, self.embed_dim_), eps=self.eps_, alloc_func=self.alloc_tensor
         )
 
@@ -42,47 +75,6 @@ def _get_qkv(
         cache_kv = layer_weight.kv_proj.mm(input).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
         return q, cache_kv
 
-    def _context_attention_kernel(
-        self, q, kv, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight, out=None
-    ) -> torch.Tensor:
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        context_attention_fwd(
-            q.view(-1, self.tp_q_head_num_, self.head_dim_),
-            kv[:, 0 : self.tp_k_head_num_, :],
-            kv[:, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :],
-            o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_),
-            infer_state.b_req_idx,
-            layer_weight.tp_alibi,
-            infer_state.b_start_loc,
-            infer_state.b_seq_len,
-            infer_state.b_ready_cache_len,
-            infer_state.max_len_in_batch,
-            infer_state.req_manager.req_to_token_indexs,
-        )
-        return o_tensor
-
-    def _token_attention_kernel(
-        self, q, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight, out=None
-    ) -> torch.Tensor:
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        token_attention_fwd(
-            q.view(-1, self.tp_q_head_num_, self.head_dim_),
-            kv[:, 0 : self.tp_k_head_num_, :],
-            kv[:, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :],
-            o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_),
-            layer_weight.tp_alibi,
-            infer_state.req_manager.req_to_token_indexs,
-            infer_state.b_req_idx,
-            infer_state.b_start_loc,
-            infer_state.b_seq_len,
-            infer_state.max_len_in_batch,
-            infer_state.total_token_num,
-            alloc_tensor_func=self.alloc_tensor,
-        )
-        return o_tensor
-
     def _get_o(self, input, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight) -> torch.Tensor:
         o_tensor = layer_weight.o_proj.mm(input.view(-1, self.tp_o_head_num_ * self.head_dim_))
         return o_tensor
diff --git a/lightllm/models/bloom/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/bloom/layer_weights/pre_and_post_layer_weight.py
index afc8c93081..000a069129 100644
--- a/lightllm/models/bloom/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/bloom/layer_weights/pre_and_post_layer_weight.py
@@ -1,24 +1,28 @@
-import torch
-import numpy as np
 from lightllm.common.basemodel import PreAndPostLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, NoTpNormWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LayerNormWeight
 
 
 class BloomPreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
-        self.pre_norm_weight_ = NoTpNormWeight(
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
+        hidden_size = network_config["hidden_size"]
+        vocab_size = network_config["vocab_size"]
+        self.pre_norm_weight_ = LayerNormWeight(
+            dim=hidden_size,
             weight_name="word_embeddings_layernorm.weight",
             data_type=self.data_type_,
             bias_name="word_embeddings_layernorm.bias",
         )
-        self.final_norm_weight_ = NoTpNormWeight(
+        self.final_norm_weight_ = LayerNormWeight(
+            dim=hidden_size,
             weight_name="ln_f.weight",
             data_type=self.data_type_,
             bias_name="ln_f.bias",
         )
 
         self.wte_weight_ = EmbeddingWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
             weight_name="word_embeddings.weight",
             data_type=self.data_type_,
         )
diff --git a/lightllm/models/bloom/layer_weights/transformer_layer_weight.py b/lightllm/models/bloom/layer_weights/transformer_layer_weight.py
index 7b27ce6f2c..568a9c0381 100644
--- a/lightllm/models/bloom/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/bloom/layer_weights/transformer_layer_weight.py
@@ -48,8 +48,8 @@ def get_slopes_power_of_2(n):
 
 
 class BloomTransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode, quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _parse_config(self):
@@ -108,18 +108,18 @@ def load_hf_weights(self, weights):
 
     def _init_ffn(self):
         self.gate_up_proj = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_inter],
             weight_names=self._gate_up_weight_name,
             data_type=self.data_type_,
             bias_names=self._gate_up_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="gate_up_proj",
+            quant_method=self.get_quant_method("gate_up_proj"),
         )
         self.down_proj = COLMMWeight(
+            in_dim=self.n_inter,
+            out_dims=[self.n_embed],
             weight_names=self._down_weight_name,
             data_type=self.data_type_,
             bias_names=self._down_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="down_proj",
+            quant_method=self.get_quant_method("down_proj"),
         )
diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py
index 7e44ec2ebf..925620bf96 100644
--- a/lightllm/models/bloom/model.py
+++ b/lightllm/models/bloom/model.py
@@ -5,6 +5,7 @@
 from lightllm.models.bloom.layer_weights.pre_and_post_layer_weight import BloomPreAndPostLayerWeight
 from lightllm.models.bloom.layer_weights.transformer_layer_weight import BloomTransformerLayerWeight
 from lightllm.common.basemodel import InferStateInfo, TpPartBaseModel
+from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend
 
 
 @ModelRegistry("bloom")
@@ -35,3 +36,8 @@ def _init_config(self):
     def _reset_num_key_value_heads(self):
         self.config["num_key_value_heads"] = self.config["num_attention_heads"]
         return
+
+    def _init_att_backend(self):
+        self.prefill_att_backend = TritonAttBackend(self)
+        self.decode_att_backend = TritonAttBackend(self)
+        return
diff --git a/lightllm/models/chatglm2/layer_infer/transformer_layer_infer.py b/lightllm/models/chatglm2/layer_infer/transformer_layer_infer.py
deleted file mode 100755
index 07ffc4beab..0000000000
--- a/lightllm/models/chatglm2/layer_infer/transformer_layer_infer.py
+++ /dev/null
@@ -1,28 +0,0 @@
-import torch
-from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
-from lightllm.models.chatglm2.layer_weights.transformer_layer_weight import ChatGLM2TransformerLayerWeight
-
-
-class ChatGLM2TransformerLayerInfer(LlamaTransformerLayerInfer):
-    """ """
-
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
-        return
-
-    def swiglu(self, x):
-        x = torch.chunk(x, 2, dim=-1)
-        return torch.nn.functional.silu(x[0]) * x[1]
-
-    def _ffn(
-        self, input, infer_state: LlamaInferStateInfo, layer_weight: ChatGLM2TransformerLayerWeight
-    ) -> torch.Tensor:
-        input = input.view(-1, self.embed_dim_)
-        up_gate_out = layer_weight.gate_up_proj.mm(input)
-        input = None
-        ffn1_out = self.swiglu(up_gate_out)
-        up_gate_out = None
-        ffn2_out = layer_weight.down_proj.mm(ffn1_out)
-        ffn1_out = None
-        return ffn2_out
diff --git a/lightllm/models/chatglm2/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/chatglm2/layer_weights/pre_and_post_layer_weight.py
deleted file mode 100644
index 0139eb8837..0000000000
--- a/lightllm/models/chatglm2/layer_weights/pre_and_post_layer_weight.py
+++ /dev/null
@@ -1,20 +0,0 @@
-from lightllm.common.basemodel import PreAndPostLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, NoTpNormWeight
-
-
-class ChatGLM2PreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
-
-        self.wte_weight_ = EmbeddingWeight(
-            weight_name="transformer.embedding.word_embeddings.weight", data_type=self.data_type_
-        )
-        self.lm_head_weight_ = LMHeadWeight(
-            weight_name="transformer.output_layer.weight",
-            data_type=self.data_type_,
-        )
-        self.final_norm_weight_ = NoTpNormWeight(
-            weight_name="transformer.encoder.final_layernorm.weight",
-            data_type=self.data_type_,
-            bias_name=None,
-        )
diff --git a/lightllm/models/chatglm2/layer_weights/transformer_layer_weight.py b/lightllm/models/chatglm2/layer_weights/transformer_layer_weight.py
deleted file mode 100755
index d4dd1b7a29..0000000000
--- a/lightllm/models/chatglm2/layer_weights/transformer_layer_weight.py
+++ /dev/null
@@ -1,72 +0,0 @@
-from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
-
-
-class ChatGLM2TransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(
-            layer_num,
-            data_type,
-            network_config,
-            mode,
-            quant_cfg,
-        )
-        return
-
-    def _preprocess_weight(self, weights):
-        n_kv_embed = self.head_dim * self.n_kv_head
-        qkv_weight_name = f"transformer.encoder.layers.{self.layer_num_}.self_attention.query_key_value.weight"
-        if qkv_weight_name in weights:
-            qkv_weight_ = weights[qkv_weight_name]
-            weights[self._q_weight_name] = qkv_weight_[: self.n_embed, :]
-            weights[self._k_weight_name] = qkv_weight_[self.n_embed : self.n_embed + n_kv_embed, :]
-            weights[self._v_weight_name] = qkv_weight_[self.n_embed + n_kv_embed : self.n_embed + 2 * n_kv_embed, :]
-            del weights[qkv_weight_name]
-
-        qkv_bias_name = f"transformer.encoder.layers.{self.layer_num_}.self_attention.query_key_value.bias"
-        if qkv_bias_name in weights:
-            qkv_bias_ = weights[qkv_bias_name]
-            weights[self._q_bias_name] = qkv_bias_[: self.n_embed]
-            weights[self._k_bias_name] = qkv_bias_[self.n_embed : self.n_embed + n_kv_embed]
-            weights[self._v_bias_name] = qkv_bias_[self.n_embed + n_kv_embed : self.n_embed + 2 * n_kv_embed]
-            del weights[qkv_bias_name]
-
-        gate_up_weight_name = f"transformer.encoder.layers.{self.layer_num_}.mlp.dense_h_to_4h.weight"
-        if gate_up_weight_name in weights:
-            gate_up_weight_ = weights[gate_up_weight_name]
-            weights[self._gate_weight_name] = gate_up_weight_[: self.n_inter, :]
-            weights[self._up_weight_name] = gate_up_weight_[self.n_inter : 2 * self.n_inter, :]
-            del weights[gate_up_weight_name]
-
-    def _parse_config(self):
-        self.n_embed = self.network_config_["hidden_size"]
-        self.n_head = self.network_config_["num_attention_heads"]
-        self.n_inter = self.network_config_["ffn_hidden_size"]
-        self.n_kv_head = self.network_config_["multi_query_group_num"]
-        self.head_dim = self.network_config_.get("head_dim", self.n_embed // self.n_head)
-
-    def load_hf_weights(self, weights):
-        self._preprocess_weight(weights)
-        super().load_hf_weights(weights)
-        return
-
-    def _init_weight_names(self):
-        self._q_weight_name = f"transformer.encoder.layers.{self.layer_num_}.self_attention.q_proj.weight"
-        self._q_bias_name = f"transformer.encoder.layers.{self.layer_num_}.self_attention.q_proj.bias"
-        self._k_weight_name = f"transformer.encoder.layers.{self.layer_num_}.self_attention.k_proj.weight"
-        self._k_bias_name = f"transformer.encoder.layers.{self.layer_num_}.self_attention.k_proj.bias"
-        self._v_weight_name = f"transformer.encoder.layers.{self.layer_num_}.self_attention.v_proj.weight"
-        self._v_bias_name = f"transformer.encoder.layers.{self.layer_num_}.self_attention.v_proj.bias"
-        self._o_weight_name = f"transformer.encoder.layers.{self.layer_num_}.self_attention.dense.weight"
-        self._o_bias_name = None
-
-        self._gate_weight_name = f"transformer.encoder.layers.{self.layer_num_}.mlp.gate_proj.weight"
-        self._gate_bias_name = None
-        self._up_weight_name = f"transformer.encoder.layers.{self.layer_num_}.mlp.up_proj.weight"
-        self._up_bias_name = None
-        self._down_weight_name = f"transformer.encoder.layers.{self.layer_num_}.mlp.dense_4h_to_h.weight"
-        self._down_bias_name = None
-
-        self._att_norm_weight_name = f"transformer.encoder.layers.{self.layer_num_}.input_layernorm.weight"
-        self._att_norm_bias_name = None
-        self._ffn_norm_weight_name = f"transformer.encoder.layers.{self.layer_num_}.post_attention_layernorm.weight"
-        self._ffn_norm_bias_name = None
diff --git a/lightllm/models/chatglm2/model.py b/lightllm/models/chatglm2/model.py
deleted file mode 100644
index e6aa395275..0000000000
--- a/lightllm/models/chatglm2/model.py
+++ /dev/null
@@ -1,78 +0,0 @@
-import os
-import json
-import torch
-
-from lightllm.models.registry import ModelRegistry
-from lightllm.models.chatglm2.layer_infer.transformer_layer_infer import ChatGLM2TransformerLayerInfer
-from lightllm.models.chatglm2.layer_weights.transformer_layer_weight import ChatGLM2TransformerLayerWeight
-from lightllm.models.chatglm2.layer_weights.pre_and_post_layer_weight import ChatGLM2PreAndPostLayerWeight
-from lightllm.models.llama.model import LlamaTpPartModel
-from lightllm.common.build_utils import repair_config
-from lightllm.utils.log_utils import init_logger
-
-logger = init_logger(__name__)
-
-
-@ModelRegistry("chatglm")
-class ChatGlm2TpPartModel(LlamaTpPartModel):
-    # Please use the fast tokenizer from:
-    # [THUDM/chatglm3-6b PR #12](https://huggingface.co/THUDM/chatglm3-6b/discussions/12).
-
-    # weight class
-    pre_and_post_weight_class = ChatGLM2PreAndPostLayerWeight
-    transformer_weight_class = ChatGLM2TransformerLayerWeight
-
-    # infer class
-    transformer_layer_infer_class = ChatGLM2TransformerLayerInfer
-
-    def __init__(self, kvargs):
-        super().__init__(kvargs)
-
-    def _init_config(self):
-        super()._init_config()
-        # rename key
-        # repair_config()
-        repair_config(self.config, same_names=["num_hidden_layers", "n_layer", "num_layers"])
-        repair_config(self.config, same_names=["vocab_size", "padded_vocab_size"])
-        repair_config(self.config, same_names=["rms_norm_eps", "layernorm_epsilon"])
-        repair_config(self.config, same_names=["seq_length", "max_sequence_length"])
-        return
-
-    def _reset_num_key_value_heads(self):
-        self.config["num_key_value_heads"] = self.config["multi_query_group_num"]
-        return
-
-    def _verify_params(self):
-        assert self.load_way == "HF", "ChatGLM only support HF format for now"
-        assert self.tp_world_size_ in [1, 2], "ChatGLM can only run in tp=1 or tp=2 for now"
-
-    def _init_to_get_rotary(self, base=10000):
-        if self.config.get("rope_scaling", {}) is None:
-            rope_scaling_factor = 1.0
-        else:
-            rope_scaling_factor = self.config.get("rope_scaling", {}).get("factor", 1.0)
-        if "max_sequence_length" in self.config:
-            max_seq_len = self.config["max_sequence_length"]
-        else:
-            max_seq_len = self.config.get("max_position_embeddings", 2048) * rope_scaling_factor
-
-        base = float(base) * self.config.get("rope_ratio", 1.0)
-
-        # NTK
-        try:
-            ntk_alpha = float(os.environ.get("LIGHTLLM_NTK_ALPHA", 1))
-            assert ntk_alpha >= 1
-            if ntk_alpha > 1:
-                logger.info(f"Note: NTK enabled, alpha set to {ntk_alpha}")
-            max_seq_len *= ntk_alpha
-            base = base * (ntk_alpha ** (self.head_dim_ / (self.head_dim_ - 2)))  # Base change formula
-        except:
-            pass
-        n_elem = self.head_dim_ // 2
-        inv_freq = 1.0 / (base ** (torch.arange(0, n_elem, 2, device="cpu", dtype=torch.float32) / n_elem))
-        t = torch.arange(max_seq_len + 1024 * 64, device="cpu", dtype=torch.float32) / rope_scaling_factor
-        freqs = torch.outer(t, inv_freq)
-
-        self._cos_cached = torch.cos(freqs).to(self.data_type).cuda()
-        self._sin_cached = torch.sin(freqs).to(self.data_type).cuda()
-        return
diff --git a/lightllm/models/chatglm2/triton_kernel/rotary_emb.py b/lightllm/models/chatglm2/triton_kernel/rotary_emb.py
deleted file mode 100755
index ad1d1c2cf0..0000000000
--- a/lightllm/models/chatglm2/triton_kernel/rotary_emb.py
+++ /dev/null
@@ -1,160 +0,0 @@
-import torch
-
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def _rotary_kernel(
-    Q,
-    K,
-    Cos,
-    Sin,
-    stride_qbs,
-    stride_qh,
-    stride_qd,
-    stride_kbs,
-    stride_kh,
-    stride_kd,
-    stride_cosbs,
-    stride_cosd,
-    stride_sinbs,
-    stride_sind,
-    max_total_len,
-    HEAD_Q,
-    HEAD_K,  # N_CTX 代表要计算的上下文长度
-    BLOCK_HEAD: tl.constexpr,
-    BLOCK_SEQ: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr,
-):
-    cur_head_index = tl.program_id(0)
-    cur_seq_index = tl.program_id(1)
-
-    cur_head_range = cur_head_index * BLOCK_HEAD + tl.arange(0, BLOCK_HEAD)
-    cur_seq_range = cur_seq_index * BLOCK_SEQ + tl.arange(0, BLOCK_SEQ)
-
-    dim_range0 = tl.arange(0, BLOCK_DMODEL // 2) * 2
-    dim_range1 = dim_range0 + 1
-
-    off_q0 = (
-        cur_seq_range[:, None, None] * stride_qbs
-        + cur_head_range[None, :, None] * stride_qh
-        + dim_range0[None, None, :] * stride_qd
-    )
-    off_q1 = (
-        cur_seq_range[:, None, None] * stride_qbs
-        + cur_head_range[None, :, None] * stride_qh
-        + dim_range1[None, None, :] * stride_qd
-    )
-
-    cos_range = tl.arange(0, BLOCK_DMODEL // 2)
-    off_dimcos_sin = cur_seq_range[:, None, None] * stride_cosbs + cos_range[None, None, :] * stride_cosd
-
-    q0 = tl.load(
-        Q + off_q0,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_Q),
-        other=0.0,
-    )
-    q1 = tl.load(
-        Q + off_q1,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_Q),
-        other=0.0,
-    )
-
-    cos = tl.load(Cos + off_dimcos_sin, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-    sin = tl.load(Sin + off_dimcos_sin, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-
-    out0 = q0 * cos - q1 * sin
-    out1 = q0 * sin + q1 * cos
-
-    tl.store(
-        Q + off_q0, out0, mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_Q)
-    )
-    tl.store(
-        Q + off_q1, out1, mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_Q)
-    )
-
-    off_k0 = (
-        cur_seq_range[:, None, None] * stride_kbs
-        + cur_head_range[None, :, None] * stride_kh
-        + dim_range0[None, None, :] * stride_kd
-    )
-    off_k1 = (
-        cur_seq_range[:, None, None] * stride_kbs
-        + cur_head_range[None, :, None] * stride_kh
-        + dim_range1[None, None, :] * stride_kd
-    )
-
-    off_dimcos_sin = cur_seq_range[:, None, None] * stride_cosbs + cos_range[None, None, :] * stride_cosd
-
-    k0 = tl.load(
-        K + off_k0,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_K),
-        other=0.0,
-    )
-    k1 = tl.load(
-        K + off_k1,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_K),
-        other=0.0,
-    )
-
-    cos = tl.load(Cos + off_dimcos_sin, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-    sin = tl.load(Sin + off_dimcos_sin, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-
-    out_k0 = k0 * cos - k1 * sin
-    out_k1 = k0 * sin + k1 * cos
-
-    tl.store(
-        K + off_k0,
-        out_k0,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_K),
-    )
-    tl.store(
-        K + off_k1,
-        out_k1,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_K),
-    )
-    return
-
-
-@torch.no_grad()
-def rotary_emb_fwd(q, k, cos, sin):
-    total_len = q.shape[0]
-    head_num_q, head_num_k = q.shape[1], k.shape[1]
-    head_dim = q.shape[2] // 2
-    assert q.shape[0] == cos.shape[0] and q.shape[0] == sin.shape[0], f"q shape {q.shape} cos shape {cos.shape}"
-    assert k.shape[0] == cos.shape[0] and k.shape[0] == sin.shape[0], f"k shape {k.shape} cos shape {cos.shape}"
-
-    BLOCK_SEQ = 16
-    BLOCK_HEAD = 4
-    if head_dim >= 128:
-        num_warps = 8
-    else:
-        num_warps = 4
-
-    grid = (triton.cdiv(head_num_q, BLOCK_HEAD), triton.cdiv(total_len, BLOCK_SEQ))
-    _rotary_kernel[grid](
-        q,
-        k,
-        cos,
-        sin,
-        q.stride(0),
-        q.stride(1),
-        q.stride(2),
-        k.stride(0),
-        k.stride(1),
-        k.stride(2),
-        cos.stride(0),
-        cos.stride(1),
-        sin.stride(0),
-        sin.stride(1),
-        total_len,
-        head_num_q,
-        head_num_k,
-        BLOCK_HEAD=BLOCK_HEAD,
-        BLOCK_SEQ=BLOCK_SEQ,
-        BLOCK_DMODEL=head_dim,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
diff --git a/lightllm/models/cohere/infer_struct.py b/lightllm/models/cohere/infer_struct.py
deleted file mode 100644
index d9571af92b..0000000000
--- a/lightllm/models/cohere/infer_struct.py
+++ /dev/null
@@ -1,8 +0,0 @@
-from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-
-
-class CohereInferStateInfo(LlamaInferStateInfo):
-    def __init__(self):
-        super().__init__()
-        self._attn_out = None
-        self._ffn_out = None
diff --git a/lightllm/models/cohere/layer_infer/post_layer_infer.py b/lightllm/models/cohere/layer_infer/post_layer_infer.py
deleted file mode 100644
index 67987a8d3b..0000000000
--- a/lightllm/models/cohere/layer_infer/post_layer_infer.py
+++ /dev/null
@@ -1,71 +0,0 @@
-import torch
-import numpy as np
-from lightllm.models.cohere.infer_struct import CohereInferStateInfo
-from lightllm.models.cohere.layer_weights.pre_and_post_layer_weight import CoherePreAndPostLayerWeight
-from lightllm.models.cohere.triton_kernels.layernorm import layernorm_forward
-from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer
-from lightllm.common.build_utils import repair_config
-from lightllm.distributed.communication_op import all_gather
-
-
-class CoherePostLayerInfer(LlamaPostLayerInfer):
-    def __init__(self, network_config, mode):
-        repair_config(config=network_config, same_names=["layer_norm_eps", "rms_norm_eps"])
-        super().__init__(network_config, mode)
-        self.eps_ = network_config["layer_norm_eps"]
-        self.logits_scale = network_config["logit_scale"]
-        return
-
-    def _norm(
-        self, input: torch.Tensor, infer_state: CohereInferStateInfo, layer_weight: CoherePreAndPostLayerWeight
-    ) -> torch.Tensor:
-        return layernorm_forward(
-            input.unsqueeze(1), layer_weight.final_norm_weight_.weight.unsqueeze(0), eps=self.eps_
-        ).squeeze(1)
-
-    def token_forward(
-        self, input_embdings: torch.Tensor, infer_state: CohereInferStateInfo, layer_weight: CoherePreAndPostLayerWeight
-    ):
-        last_input, token_num = self._slice_get_last_input(input_embdings, infer_state)
-        input_embdings_dtype = input_embdings.dtype
-        input_embdings = None
-        last_input = self._norm(last_input, infer_state, layer_weight)
-        last_input = last_input.permute(1, 0).view(-1, token_num)
-        logic_batch = layer_weight.lm_head_weight_.lm_head(input=last_input, alloc_func=self.alloc_tensor)
-        last_input = None
-        vocab_size = layer_weight.lm_head_weight_.vocab_size
-        if self.tp_world_size_ == 1:
-            gather_data = logic_batch
-        else:
-            gather_data = self.alloc_tensor((vocab_size, token_num), dtype=input_embdings_dtype)
-            split_indexes = np.linspace(0, vocab_size, self.tp_world_size_ + 1, dtype=np.int64)
-            all_gather(
-                [gather_data[split_indexes[i] : split_indexes[i + 1], :] for i in range(self.tp_world_size_)],
-                logic_batch,
-                group=infer_state.dist_group,
-                async_op=False,
-            )
-        gather_data = gather_data * self.logits_scale
-        logic_batch = None
-        ans_logics = self.alloc_tensor(
-            (token_num, vocab_size),
-            dtype=torch.float32,
-        )
-        ans_logics[:, :] = gather_data.permute(1, 0)
-        gather_data = None
-        return ans_logics
-
-    def tpsp_token_forward(
-        self, input_embdings: torch.Tensor, infer_state: CohereInferStateInfo, layer_weight: CoherePreAndPostLayerWeight
-    ):
-        raise NotImplementedError("not impl")
-
-    def overlap_tpsp_token_forward(
-        self,
-        input_embdings: torch.Tensor,
-        input_embdings1: torch.Tensor,
-        infer_state: CohereInferStateInfo,
-        infer_state1: CohereInferStateInfo,
-        layer_weight: CoherePreAndPostLayerWeight,
-    ):
-        raise NotImplementedError("not impl")
diff --git a/lightllm/models/cohere/layer_infer/transformer_layer_infer.py b/lightllm/models/cohere/layer_infer/transformer_layer_infer.py
deleted file mode 100644
index 0cdd281a37..0000000000
--- a/lightllm/models/cohere/layer_infer/transformer_layer_infer.py
+++ /dev/null
@@ -1,84 +0,0 @@
-import torch
-from functools import partial
-
-from lightllm.common.basemodel.layer_infer.template.transformer_layer_infer_cohere_template import (
-    TransformerLayerCohereInferTpl,
-)
-from lightllm.models.cohere.infer_struct import CohereInferStateInfo
-from lightllm.models.cohere.layer_weights.transformer_layer_weight import CohereTransformerLayerWeight
-from lightllm.models.cohere.triton_kernels.layernorm import layernorm_forward, torch_layernorm
-from lightllm.models.cohere.triton_kernels.rotary_emb import rotary_emb_fwd
-from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
-from lightllm.models.llama.triton_kernel.silu_and_mul import silu_and_mul_fwd
-
-
-class CohereTransformerLayerInfer(TransformerLayerCohereInferTpl):
-    def __init__(self, layer_num, network_config, mode):
-        super().__init__(layer_num, network_config, mode)
-        self.tp_q_head_num_ = network_config["num_attention_heads"] // self.tp_world_size_
-        self.tp_k_head_num_ = network_config["num_key_value_heads"] // self.tp_world_size_
-        self.tp_v_head_num_ = network_config["num_key_value_heads"] // self.tp_world_size_
-        self.tp_o_head_num_ = self.tp_q_head_num_
-        self.head_dim_ = network_config["hidden_size"] // network_config["num_attention_heads"]
-        self.embed_dim_ = network_config["hidden_size"]
-        self.eps_ = self.network_config_["layer_norm_eps"]
-        self.use_qk_norm_ = network_config.get("use_qk_norm", False)
-        self._bind_func()
-
-    def _bind_func(self):
-        self._bind_rotary_emb_fwd()
-        self._bind_norm()
-        self._bind_attn()
-
-    def _rotary_emb_fwd(self, q, kv, position_cos, position_sin):
-        return rotary_emb_fwd(
-            q.view(-1, self.tp_q_head_num_, self.head_dim_),
-            kv,
-            position_cos,
-            position_sin,
-        )
-
-    def _bind_rotary_emb_fwd(self):
-        self._rotary_emb_fwd = partial(CohereTransformerLayerInfer._rotary_emb_fwd, self)
-
-    def _att_norm(self, input, infer_state, layer_weight: CohereTransformerLayerWeight):
-        return layernorm_forward(
-            input.unsqueeze(1), layer_weight.att_norm_weight_.weight.unsqueeze(0), self.eps_
-        ).squeeze(1)
-
-    def _q_norm(self, input, infer_state, layer_weight: CohereTransformerLayerWeight):
-        return layernorm_forward(input, layer_weight.q_norm_weight_.weight, self.eps_)
-
-    def _k_norm(self, input, infer_state, layer_weight: CohereTransformerLayerWeight):
-        return layernorm_forward(input, layer_weight.k_norm_weight_.weight, self.eps_)
-
-    def _bind_norm(self):
-        self._att_norm = partial(CohereTransformerLayerInfer._att_norm, self)
-        self._q_norm = partial(CohereTransformerLayerInfer._q_norm, self)
-        self._k_norm = partial(CohereTransformerLayerInfer._k_norm, self)
-
-    def _bind_attn(self):
-        # no need to re-impl
-        LlamaTransformerLayerInfer._bind_attention(self)
-
-    def _get_o(
-        self, input, infer_state: CohereInferStateInfo, layer_weight: CohereTransformerLayerWeight
-    ) -> torch.Tensor:
-        input = input.view(-1, self.tp_o_head_num_ * self.head_dim_)
-        # o_tensor = layer_weight.mm_op.apply(input, layer_weight.o_weight_)
-        o_tensor = layer_weight.o_proj.mm(input)
-        return o_tensor
-
-    def _ffn(
-        self, input, infer_state: CohereInferStateInfo, layer_weight: CohereTransformerLayerWeight
-    ) -> torch.Tensor:
-        input = input.view(-1, self.embed_dim_)
-        up_gate_out = layer_weight.gate_up_proj.mm(input)
-        ffn1_out = self.alloc_tensor((input.size(0), up_gate_out.size(1) // 2), input.dtype)
-        silu_and_mul_fwd(up_gate_out, ffn1_out)
-        input = None
-        up_gate_out = None
-        # ffn2_out = layer_weight.mm_op.apply(ffn1_out, layer_weight.down_proj)
-        ffn2_out = layer_weight.down_proj.mm(ffn1_out)
-        ffn1_out = None
-        return ffn2_out
diff --git a/lightllm/models/cohere/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/cohere/layer_weights/pre_and_post_layer_weight.py
deleted file mode 100644
index f2e5f85472..0000000000
--- a/lightllm/models/cohere/layer_weights/pre_and_post_layer_weight.py
+++ /dev/null
@@ -1,25 +0,0 @@
-from lightllm.common.basemodel import PreAndPostLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, NoTpNormWeight
-
-
-class CoherePreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
-        tie_weight = self.network_config_.get("tie_word_embeddings", True)
-
-        self.wte_weight_ = EmbeddingWeight(
-            weight_name="model.embed_tokens.weight",
-            data_type=self.data_type_,
-        )
-        if tie_weight:
-            self.lm_head_weight_ = self.wte_weight_
-        else:
-            self.lm_head_weight_ = LMHeadWeight(
-                weight_name="model.lm_head.weight",
-                data_type=self.data_type_,
-            )
-        self.final_norm_weight_ = NoTpNormWeight(
-            weight_name="model.norm.weight",
-            data_type=self.data_type_,
-            bias_name=None,
-        )
diff --git a/lightllm/models/cohere/layer_weights/transformer_layer_weight.py b/lightllm/models/cohere/layer_weights/transformer_layer_weight.py
deleted file mode 100644
index 9c446b49e9..0000000000
--- a/lightllm/models/cohere/layer_weights/transformer_layer_weight.py
+++ /dev/null
@@ -1,25 +0,0 @@
-from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import NoTpNormWeight, TpHeadNormWeight
-
-
-class CohereTransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
-        return
-
-    def _parse_config(self):
-        super()._parse_config()
-        self.use_qk_norm = self.network_config_.get("use_qk_norm", False)
-
-    def _init_norm(self):
-        self.att_norm_weight_ = NoTpNormWeight(self._att_norm_weight_name, self.data_type_)
-
-        if self.use_qk_norm:
-            self.q_norm_weight_ = TpHeadNormWeight(
-                f"model.layers.{self.layer_num_}.self_attn.q_norm.weight", self.data_type_
-            )
-            self.k_norm_weight_ = TpHeadNormWeight(
-                f"model.layers.{self.layer_num_}.self_attn.k_norm.weight", self.data_type_
-            )
-
-        return
diff --git a/lightllm/models/cohere/model.py b/lightllm/models/cohere/model.py
deleted file mode 100644
index 5b317c1331..0000000000
--- a/lightllm/models/cohere/model.py
+++ /dev/null
@@ -1,69 +0,0 @@
-import os
-import torch
-from lightllm.common.basemodel.basemodel import TpPartBaseModel
-from lightllm.common.basemodel.layer_infer.template.transformer_layer_infer_cohere_template import (
-    TransformerLayerCohereInferTpl,
-)
-from lightllm.common.kv_cache_mem_manager import MemoryManager
-from lightllm.models.registry import ModelRegistry
-from lightllm.models.cohere.infer_struct import CohereInferStateInfo
-from lightllm.models.cohere.layer_infer.post_layer_infer import CoherePostLayerInfer
-from lightllm.models.cohere.layer_infer.transformer_layer_infer import CohereTransformerLayerInfer
-from lightllm.models.cohere.layer_weights.pre_and_post_layer_weight import CoherePreAndPostLayerWeight
-from lightllm.models.cohere.layer_weights.transformer_layer_weight import CohereTransformerLayerWeight
-from lightllm.models.llama.layer_infer.pre_layer_infer import LlamaPreLayerInfer
-from lightllm.models.llama.model import LlamaTpPartModel
-from lightllm.utils.log_utils import init_logger
-
-logger = init_logger(__name__)
-
-
-@ModelRegistry("cohere")
-class CohereTpPartModel(LlamaTpPartModel):
-    pre_and_post_weight_class = CoherePreAndPostLayerWeight
-    transformer_weight_class = CohereTransformerLayerWeight
-
-    pre_layer_infer_class = LlamaPreLayerInfer
-    transformer_layer_infer_class = CohereTransformerLayerInfer
-    post_layer_infer_class = CoherePostLayerInfer
-
-    infer_state_class = CohereInferStateInfo
-
-    def _init_to_get_rotary(self, default_base=10000):
-        partial_head_dim = int(self.config.get("partial_rotary_factor", 1) * self.head_dim_)
-        if self.config.get("rope_scaling", {}) is None:
-            rope_scaling_factor = 1.0
-        else:
-            rope_scaling_factor = self.config.get("rope_scaling", {}).get("factor", 1.0)
-
-        base = self.config.get("rope_theta", float(default_base))
-
-        if "max_sequence_length" in self.config:
-            max_seq_len = self.config["max_sequence_length"]
-        else:
-            max_position_embeddings = self.config.get(
-                "max_position_embeddings", 2048 if base <= 10000.0 + 1e-5 else 16384
-            )
-            max_seq_len = max_position_embeddings * rope_scaling_factor
-
-        # NTK
-        try:
-            ntk_alpha = float(os.environ.get("LIGHTLLM_NTK_ALPHA", 1))
-            assert ntk_alpha >= 1
-            if ntk_alpha > 1:
-                logger.info(f"Note: NTK enabled, alpha set to {ntk_alpha}")
-            max_seq_len *= ntk_alpha
-            base = base * (ntk_alpha ** (partial_head_dim / (partial_head_dim - 2)))  # Base change formula
-        except:
-            pass
-
-        inv_freq = 1.0 / (
-            base ** (torch.arange(0, partial_head_dim, 2, device="cpu", dtype=torch.float32) / partial_head_dim)
-        )
-        t = torch.arange(max_seq_len + 1024 * 128, device="cpu", dtype=torch.float32) / rope_scaling_factor
-        freqs = torch.outer(t, inv_freq)
-        freqs = torch.repeat_interleave(freqs, 2, dim=-1)
-
-        self._cos_cached = torch.cos(freqs).to(self.data_type).cuda()
-        self._sin_cached = torch.sin(freqs).to(self.data_type).cuda()
-        return
diff --git a/lightllm/models/cohere/triton_kernels/layernorm.py b/lightllm/models/cohere/triton_kernels/layernorm.py
deleted file mode 100644
index c1d5ff4cd6..0000000000
--- a/lightllm/models/cohere/triton_kernels/layernorm.py
+++ /dev/null
@@ -1,131 +0,0 @@
-import torch
-import triton
-import triton.language as tl
-
-# LayerNorm adapted from triton tutorial, used for Cohere q, k norm
-# X [N, head_num, head_dim]
-# W [head_num, head_dim]
-@triton.jit
-def _layer_norm_fwd_kernel(
-    X,  # pointer to the input
-    W,  # pointer to the weights
-    Y,
-    stride_x_N,
-    stride_x_hn,
-    stride_x_hd,
-    stride_y_N,
-    stride_y_hn,
-    stride_y_hd,
-    stride_w_hn,
-    stride_w_hd,
-    N,  # number of columns in X
-    eps,  # epsilon to avoid division by zero
-    BLOCK_SIZE: tl.constexpr,
-):
-    Seq = tl.program_id(0)
-    H = tl.program_id(1)
-
-    X += Seq * stride_x_N + H * stride_x_hn
-    Y += Seq * stride_y_N + H * stride_y_hn
-    W += H * stride_w_hn
-
-    _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
-    for off in range(0, N, BLOCK_SIZE):
-        cols = off + tl.arange(0, BLOCK_SIZE)
-        a = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)
-        _mean += a
-    mean = tl.sum(_mean, axis=0) / N
-
-    _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
-    for off in range(0, N, BLOCK_SIZE):
-        cols = off + tl.arange(0, BLOCK_SIZE)
-        x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)
-        x = tl.where(cols < N, x - mean, 0.0)
-        _var += x * x
-    var = tl.sum(_var, axis=0) / N
-    rstd = 1 / tl.sqrt(var + eps)
-
-    for off in range(0, N, BLOCK_SIZE):
-        cols = off + tl.arange(0, BLOCK_SIZE)
-        mask = cols < N
-        w = tl.load(W + cols, mask=mask).to(tl.float32)
-        x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32)
-        x_hat = (x - mean) * rstd
-        y = x_hat * w
-
-        tl.store(Y + cols, y.to(X.dtype.element_ty), mask=mask)
-
-
-def layernorm_forward(
-    X,  # pointer to the input
-    W,  # pointer to the weights
-    eps,  # epsilon to avoid division by zero
-):
-    assert len(X.shape) == 3
-    assert len(W.shape) == 2
-    assert X.shape[-1] == W.shape[-1]
-    assert X.shape[-2] == W.shape[-2]
-
-    y = torch.empty_like(X)
-
-    stride_x_N = X.stride(0)
-    stride_x_hn = X.stride(1)
-    stride_x_hd = X.stride(2)
-
-    stride_y_N = y.stride(0)
-    stride_y_hn = y.stride(1)
-    stride_y_hd = y.stride(2)
-
-    stride_w_hn = W.stride(0)
-    stride_w_hd = W.stride(1)
-
-    N = X.shape[-1]
-    BLOCK_SIZE = 128
-
-    grid = (X.shape[0], X.shape[1])
-    _layer_norm_fwd_kernel[grid](
-        X,
-        W,
-        y,
-        stride_x_N,
-        stride_x_hn,
-        stride_x_hd,
-        stride_y_N,
-        stride_y_hn,
-        stride_y_hd,
-        stride_w_hn,
-        stride_w_hd,
-        N,
-        eps,
-        BLOCK_SIZE,
-    )
-
-    return y
-
-
-def torch_layernorm(x, weight, eps):
-    inp_dtype = x.dtype
-    x = x.to(torch.float32)
-    mean = x.mean(-1, keepdim=True)
-    variance = (x - mean).pow(2).mean(-1, keepdim=True)
-    x = (x - mean) * torch.rsqrt(variance + eps)
-    x = weight.to(torch.float32) * x
-    return x.to(inp_dtype)
-
-
-def test_layernorm(eps=1e-5):
-    # create data
-    dtype = torch.float16
-    x_shape = (5, 1, 128)
-    w_shape = (x_shape[-2], x_shape[-1])
-    weight = torch.rand(w_shape, dtype=dtype, device="cuda")
-    x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device="cuda")
-    # forward pass
-    y_ref = torch_layernorm(x, weight, eps).to(dtype)
-    y_out = layernorm_forward(x, weight, eps)
-
-    # compare
-    print("type:", y_out.dtype, y_ref.dtype)
-    print("max delta:", torch.max(torch.abs(y_out - y_ref)))
-    assert torch.allclose(y_out, y_ref, atol=1e-2, rtol=0)
-    return
diff --git a/lightllm/models/cohere/triton_kernels/rotary_emb.py b/lightllm/models/cohere/triton_kernels/rotary_emb.py
deleted file mode 100644
index ac338e71ef..0000000000
--- a/lightllm/models/cohere/triton_kernels/rotary_emb.py
+++ /dev/null
@@ -1,199 +0,0 @@
-import torch
-
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def _rotary_kernel(
-    Q,
-    K,
-    Cos,
-    Sin,
-    stride_qbs,
-    stride_qh,
-    stride_qd,
-    stride_kbs,
-    stride_kh,
-    stride_kd,
-    stride_cosbs,
-    stride_cosd,
-    stride_sinbs,
-    stride_sind,
-    max_total_len,
-    HEAD_Q,
-    HEAD_K,  # N_CTX 代表要计算的上下文长度
-    BLOCK_HEAD: tl.constexpr,
-    BLOCK_SEQ: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr,
-):
-    cur_head_index = tl.program_id(0)
-    cur_seq_index = tl.program_id(1)
-
-    cur_head_range = cur_head_index * BLOCK_HEAD + tl.arange(0, BLOCK_HEAD)
-    cur_seq_range = cur_seq_index * BLOCK_SEQ + tl.arange(0, BLOCK_SEQ)
-
-    dim_range0 = tl.arange(0, BLOCK_DMODEL // 2) * 2
-    dim_range1 = tl.arange(0, BLOCK_DMODEL // 2) * 2 + 1
-
-    off_q0 = (
-        cur_seq_range[:, None, None] * stride_qbs
-        + cur_head_range[None, :, None] * stride_qh
-        + dim_range0[None, None, :] * stride_qd
-    )
-    off_q1 = (
-        cur_seq_range[:, None, None] * stride_qbs
-        + cur_head_range[None, :, None] * stride_qh
-        + dim_range1[None, None, :] * stride_qd
-    )
-
-    off_dimcos_sin0 = cur_seq_range[:, None, None] * stride_cosbs + dim_range0[None, None, :] * stride_cosd
-    off_dimcos_sin1 = cur_seq_range[:, None, None] * stride_cosbs + dim_range1[None, None, :] * stride_cosd
-
-    q0 = tl.load(
-        Q + off_q0,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_Q),
-        other=0.0,
-    )
-    q1 = tl.load(
-        Q + off_q1,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_Q),
-        other=0.0,
-    )
-
-    cos0 = tl.load(Cos + off_dimcos_sin0, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-    sin0 = tl.load(Sin + off_dimcos_sin0, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-
-    cos1 = tl.load(Cos + off_dimcos_sin1, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-    sin1 = tl.load(Sin + off_dimcos_sin1, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-
-    out0 = q0 * cos0 - q1 * sin0
-    out1 = q0 * sin1 + q1 * cos1
-
-    tl.store(
-        Q + off_q0, out0, mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_Q)
-    )
-    tl.store(
-        Q + off_q1, out1, mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_Q)
-    )
-
-    off_k0 = (
-        cur_seq_range[:, None, None] * stride_kbs
-        + cur_head_range[None, :, None] * stride_kh
-        + dim_range0[None, None, :] * stride_kd
-    )
-    off_k1 = (
-        cur_seq_range[:, None, None] * stride_kbs
-        + cur_head_range[None, :, None] * stride_kh
-        + dim_range1[None, None, :] * stride_kd
-    )
-
-    off_dimcos_sin0 = cur_seq_range[:, None, None] * stride_cosbs + dim_range0[None, None, :] * stride_cosd
-    off_dimcos_sin1 = cur_seq_range[:, None, None] * stride_cosbs + dim_range1[None, None, :] * stride_cosd
-
-    k0 = tl.load(
-        K + off_k0,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_K),
-        other=0.0,
-    )
-    k1 = tl.load(
-        K + off_k1,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_K),
-        other=0.0,
-    )
-
-    cos0 = tl.load(Cos + off_dimcos_sin0, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-    sin0 = tl.load(Sin + off_dimcos_sin0, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-
-    cos1 = tl.load(Cos + off_dimcos_sin1, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-    sin1 = tl.load(Sin + off_dimcos_sin1, mask=cur_seq_range[:, None, None] < max_total_len, other=0.0)
-
-    out_k0 = k0 * cos0 - k1 * sin0
-    out_k1 = k0 * sin1 + k1 * cos1
-
-    tl.store(
-        K + off_k0,
-        out_k0,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_K),
-    )
-    tl.store(
-        K + off_k1,
-        out_k1,
-        mask=(cur_seq_range[:, None, None] < max_total_len) & (cur_head_range[None, :, None] < HEAD_K),
-    )
-    return
-
-
-def torch_cohere_rotary_emb(x, cos, sin):
-    dtype = x.dtype
-    seq_len, h, dim = x.shape
-    x = x.float()
-    x1 = x[:, :, ::2]
-    x2 = x[:, :, 1::2]
-    rot_x = torch.stack([-x2, x1], dim=-1).flatten(-2)
-    cos = cos.view((seq_len, 1, dim))
-    sin = sin.view((seq_len, 1, dim))
-    o = (x * cos) + (rot_x * sin)
-    return o.to(dtype=dtype)
-
-
-@torch.no_grad()
-def rotary_emb_fwd(q, k, cos, sin, partial_rotary_factor=1.0):
-    total_len = q.shape[0]
-    head_num_q, head_num_k = q.shape[1], k.shape[1]
-    head_dim = int(q.shape[2] * partial_rotary_factor)
-    assert q.shape[0] == cos.shape[0] and q.shape[0] == sin.shape[0], f"q shape {q.shape} cos shape {cos.shape}"
-    assert k.shape[0] == cos.shape[0] and k.shape[0] == sin.shape[0], f"k shape {k.shape} cos shape {cos.shape}"
-
-    BLOCK_SEQ = 16
-    BLOCK_HEAD = 4
-    if head_dim >= 128:
-        num_warps = 8
-    else:
-        num_warps = 4
-
-    grid = (triton.cdiv(head_num_q, BLOCK_HEAD), triton.cdiv(total_len, BLOCK_SEQ))
-    _rotary_kernel[grid](
-        q,
-        k,
-        cos,
-        sin,
-        q.stride(0),
-        q.stride(1),
-        q.stride(2),
-        k.stride(0),
-        k.stride(1),
-        k.stride(2),
-        cos.stride(0),
-        cos.stride(1),
-        sin.stride(0),
-        sin.stride(1),
-        total_len,
-        head_num_q,
-        head_num_k,
-        BLOCK_HEAD=BLOCK_HEAD,
-        BLOCK_SEQ=BLOCK_SEQ,
-        BLOCK_DMODEL=head_dim,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
-
-
-def test_rotary_emb(SEQ_LEN, H, D, dtype, eps=1e-5, device="cuda"):
-    # create data
-    x_shape = (SEQ_LEN, H, D)
-    x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device="cuda")
-    y = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device="cuda")
-    cos_shape = (SEQ_LEN, D)
-    cos = -1.2 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="cuda")
-    sin = -2.0 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="cuda")
-    # forward pass
-    y_tri = torch_cohere_rotary_emb(x, cos, sin)
-    rotary_emb_fwd(x, y, cos, sin)
-    y_ref = x
-
-    # compare
-    print("type:", y_tri.dtype, y_ref.dtype)
-    print("max delta:", torch.max(torch.abs(y_tri - y_ref)))
-    assert torch.allclose(y_tri, y_ref, atol=1e-2, rtol=0)
diff --git a/lightllm/models/deepseek2/flashattention_infer_struct.py b/lightllm/models/deepseek2/flashattention_infer_struct.py
deleted file mode 100644
index 72ba8a43b1..0000000000
--- a/lightllm/models/deepseek2/flashattention_infer_struct.py
+++ /dev/null
@@ -1,65 +0,0 @@
-import os
-import torch
-import numpy as np
-import torch.distributed as dist
-from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
-from lightllm.utils.dist_utils import get_current_device_id
-from lightllm.utils.envs_utils import get_env_start_args
-from lightllm.common.basemodel.triton_kernel.fa3_utils import page_table_copy
-
-
-class Deepseek2FlashAttentionStateInfo(Deepseek2InferStateInfo):
-    _shared_page_table_buffer = None
-
-    def __init__(self):
-        super().__init__()
-
-    @classmethod
-    def get_page_table_buffer(cls, graph_max_batch_size: int, max_seq_len: int):
-        if cls._shared_page_table_buffer is None:
-            cls._shared_page_table_buffer = [
-                torch.empty(graph_max_batch_size * max_seq_len, dtype=torch.int32).to(get_current_device_id()),
-                torch.empty(graph_max_batch_size * max_seq_len, dtype=torch.int32).to(get_current_device_id()),
-            ]
-        return cls._shared_page_table_buffer
-
-    def init_some_extra_state(self, model):
-        super().init_some_extra_state(model)
-        args_mtp_step = get_env_start_args().mtp_step
-        if self.is_prefill:
-            self.cu_seqlens_q = self.b1_cu_q_seq_len
-            self.cu_seqlens_k = self.b1_cu_kv_seq_len
-            self.has_prefix_kv = self.max_cache_len > 0
-            if self.has_prefix_kv:
-                self.cu_seqlens_prefix_k = torch.nn.functional.pad(
-                    torch.cumsum(self.b_ready_cache_len, dim=0, dtype=torch.int32), (1, 0)
-                )
-                self.prefix_k_max_len = self.max_cache_len
-                self.prefix_total_token_num = self.prefix_total_token_num
-        else:
-            # Meta information of flashattention for decoding
-            self.cu_seqlens_q = self.b1_cu_q_seq_len
-            self.cu_seqlens_k = self.b1_cu_kv_seq_len
-            max_seq_len_k = self.max_kv_seq_len
-            att_batch_size = self.batch_size // (args_mtp_step + 1)
-            if self.batch_size <= model.graph_max_batch_size and self.max_len_in_batch <= model.graph_max_len_in_batch:
-                page_buffer = Deepseek2FlashAttentionStateInfo.get_page_table_buffer(
-                    model.graph_max_batch_size, model.graph_max_len_in_batch
-                )
-                self.page_table = page_buffer[self.microbatch_index][
-                    : att_batch_size * model.graph_max_len_in_batch
-                ].view(att_batch_size, model.graph_max_len_in_batch)
-            else:
-                self.page_table = torch.empty((att_batch_size, self.max_len_in_batch), dtype=torch.int32).to(
-                    self.input_ids.device
-                )
-            page_table_copy(
-                page_table=self.page_table[:, :max_seq_len_k],
-                req_to_token_indexs=model.req_manager.req_to_token_indexs,
-                b_req_idx=self.b_req_idx[args_mtp_step :: (args_mtp_step + 1)],
-            )
-            if args_mtp_step > 0:
-                self.b_att_seq_len = self.b_seq_len[args_mtp_step :: (args_mtp_step + 1)].contiguous()
-            else:
-                self.b_att_seq_len = self.b_seq_len
-        return
diff --git a/lightllm/models/deepseek2/flashinfer_struct.py b/lightllm/models/deepseek2/flashinfer_struct.py
deleted file mode 100644
index db6386f797..0000000000
--- a/lightllm/models/deepseek2/flashinfer_struct.py
+++ /dev/null
@@ -1,106 +0,0 @@
-import os
-import torch
-import numpy as np
-import torch.distributed as dist
-from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
-from lightllm.utils.envs_utils import get_env_start_args
-from lightllm.models.deepseek2.triton_kernel.repack_kv_index import repack_kv_index
-
-
-class Deepseek2FlashInferStateInfo(Deepseek2InferStateInfo):
-    def __init__(self):
-        super().__init__()
-        self.prefill_wrapper = None
-        self.decode_wrapper = None
-        self.flashinfer_extra_state = None
-
-    def init_some_extra_state(self, model):
-        super().init_some_extra_state(model)
-        self.flashinfer_extra_state = model.flashinfer_extra_state
-
-        import flashinfer
-
-        if not self.is_prefill:
-            if get_env_start_args().enable_flashinfer_decode:
-                self.q_indptr = torch.arange(self.batch_size + 1, dtype=torch.int32).to(self.input_ids.device)
-                if self.batch_size <= model.graph_max_batch_size:
-                    self.kv_indices = self.flashinfer_extra_state.kv_indices_buffer[self.microbatch_index][
-                        : self.batch_size * self.flashinfer_extra_state.max_seq_length
-                    ]
-                else:
-                    self.kv_indices = torch.empty(
-                        self.batch_size * self.flashinfer_extra_state.max_seq_length,
-                        dtype=torch.int32,
-                        device=self.input_ids.device,
-                    )
-                repack_kv_index(
-                    self.req_manager.req_to_token_indexs,
-                    self.b_req_idx,
-                    self.b_seq_len,
-                    self.b_start_loc,
-                    self.max_len_in_batch,
-                    self.kv_indices,
-                )
-                if self.decode_wrapper is None:
-                    self.decode_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
-                        self.flashinfer_extra_state.workspace_buffer,
-                        use_cuda_graph=True,
-                        qo_indptr=self.q_indptr,
-                        kv_indices=self.kv_indices,
-                        kv_indptr=self.kv_starts,
-                        kv_len_arr=self.b_seq_len,
-                    )
-                    self.decode_wrapper.plan(
-                        self.q_indptr,
-                        self.kv_starts,
-                        self.kv_indices,
-                        self.b_seq_len,
-                        self.flashinfer_extra_state.tp_q_head_num,
-                        self.flashinfer_extra_state.kv_lora_rank,
-                        self.flashinfer_extra_state.qk_rope_head_dim,
-                        1,
-                        False,  # causal
-                        self.flashinfer_extra_state.softmax_scale,
-                        self.flashinfer_extra_state.q_data_type,
-                        self.flashinfer_extra_state.kv_data_type,
-                    )
-        else:
-            if get_env_start_args().enable_flashinfer_prefill:
-                q_starts = self.b1_cu_q_seq_len.int()
-                kv_starts = self.b1_kv_start_loc.int()
-                if self.prefill_wrapper is None:
-                    self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper(
-                        self.flashinfer_extra_state.workspace_buffer, "NHD"
-                    )
-                self.prefill_wrapper.plan(
-                    qo_indptr=q_starts,
-                    kv_indptr=kv_starts,
-                    num_qo_heads=self.flashinfer_extra_state.tp_q_head_num,
-                    num_kv_heads=self.flashinfer_extra_state.tp_q_head_num,
-                    head_dim_qk=self.flashinfer_extra_state.qk_nope_head_dim
-                    + self.flashinfer_extra_state.qk_rope_head_dim,
-                    head_dim_vo=self.flashinfer_extra_state.qk_nope_head_dim,
-                    q_data_type=self.flashinfer_extra_state.q_data_type,
-                    causal=True,
-                    sm_scale=self.flashinfer_extra_state.softmax_scale,
-                )
-        return
-
-    def copy_for_cuda_graph(self, new_infer_state):
-        super().copy_for_cuda_graph(new_infer_state)
-        if get_env_start_args().enable_flashinfer_decode and not self.is_prefill:
-            self.decode_wrapper.plan(
-                new_infer_state.q_indptr,
-                new_infer_state.kv_starts,
-                new_infer_state.kv_indices,
-                new_infer_state.b_seq_len,
-                new_infer_state.flashinfer_extra_state.tp_q_head_num,
-                new_infer_state.flashinfer_extra_state.kv_lora_rank,
-                new_infer_state.flashinfer_extra_state.qk_rope_head_dim,
-                1,
-                False,  # causal
-                new_infer_state.flashinfer_extra_state.softmax_scale,
-                new_infer_state.flashinfer_extra_state.q_data_type,
-                new_infer_state.flashinfer_extra_state.kv_data_type,
-            )
-        return
diff --git a/lightllm/models/deepseek2/infer_struct.py b/lightllm/models/deepseek2/infer_struct.py
index 0c2ef30489..4dd79305c3 100644
--- a/lightllm/models/deepseek2/infer_struct.py
+++ b/lightllm/models/deepseek2/infer_struct.py
@@ -1,21 +1,6 @@
-import os
-import torch
-import numpy as np
-import torch.distributed as dist
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
 
 
 class Deepseek2InferStateInfo(LlamaInferStateInfo):
     def __init__(self):
         super().__init__()
-        self.kv_starts = None
-
-    def init_some_extra_state(self, model):
-        super().init_some_extra_state(model)
-        if not self.is_prefill:
-            self.kv_starts = self.b1_cu_kv_seq_len
-
-        if self.is_prefill:
-            self.b1_kv_start_loc = self.b1_cu_kv_seq_len
-            self.max_value_in_b_seq_len = self.max_kv_seq_len
-        return
diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py
index ff20bc6ee6..e1e435cce4 100644
--- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py
@@ -1,40 +1,25 @@
 import os
 import torch
-import torch.functional as F
 import torch.distributed as dist
-import numpy as np
 import triton
-from typing import Tuple
 from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import Deepseek2TransformerLayerWeight
-from lightllm.models.deepseek2.triton_kernel.destindex_copy_kv import destindex_copy_kv
-from lightllm.models.deepseek2.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8
-from lightllm.models.deepseek2.triton_kernel.context_flashattention_nopad import (
-    context_attention_fwd,
-)
-from lightllm.models.deepseek2.triton_kernel.context_flashattention_nopad_fp8 import context_attention_fwd_fp8
-from lightllm.models.deepseek2.triton_kernel.context_flashattention_nopad_with_v import context_attention_fwd_with_v
+from lightllm.common.basemodel.attention.base_att import AttControl
 from lightllm.models.deepseek2.triton_kernel.sample_kv import sample_kv
-from lightllm.models.deepseek2.triton_kernel.repeat_rope import repeat_rope
-from lightllm.models.deepseek2.triton_kernel.gqa_flash_decoding import gqa_token_decode_attention_flash_decoding
-from lightllm.models.deepseek2.triton_kernel.gqa_flash_decoding_fp8 import gqa_token_decode_attention_flash_decoding_fp8
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
 from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd
 from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
-from lightllm.models.deepseek2.flashinfer_struct import Deepseek2FlashInferStateInfo
-from lightllm.models.deepseek2.flashattention_infer_struct import Deepseek2FlashAttentionStateInfo
 from functools import partial
 from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale
 from lightllm.distributed.communication_op import all_gather, all_gather_into_tensor, all_reduce, reduce_scatter_tensor
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.dist_utils import get_global_world_size
 from lightllm.utils.log_utils import init_logger
-from lightllm.utils.sgl_utils import flash_attn_varlen_func, flash_attn_with_kvcache, merge_state_v2
 
 logger = init_logger(__name__)
 
 
 class Deepseek2TransformerLayerInfer(LlamaTransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
+    def __init__(self, layer_num, network_config):
         self.tp_k_head_num_ = 1
         self.tp_v_head_num_ = 1
         self.qk_nope_head_dim = network_config["qk_nope_head_dim"]
@@ -66,7 +51,7 @@ def __init__(self, layer_num, network_config, mode=[]):
                 mscale = get_deepseek_mscale(scaling_factor, mscale_all_dim)
                 self.softmax_scale = self.softmax_scale * mscale * mscale
         self.enable_cc_method = not os.getenv("DISABLE_CC_METHOD", "False").upper() in ["ON", "TRUE", "1"]
-        super().__init__(layer_num, network_config, mode)
+        super().__init__(layer_num, network_config)
         self.num_heads = network_config["num_attention_heads"]
         self.num_kv_heads = network_config["num_key_value_heads"]
         return
@@ -78,8 +63,8 @@ def _bind_func(self):
 
     def _bind_ffn(self):
         if self.is_moe:
-            moe_mode = os.environ.get("MOE_MODE", "TP")
-            if moe_mode == "EP":
+            enable_ep_moe = get_env_start_args().enable_ep_moe
+            if enable_ep_moe:
                 self._ffn = partial(Deepseek2TransformerLayerInfer._moe_ffn_edp, self)
                 self._tpsp_ffn = self._tpsp_ffn_ep
             else:
@@ -89,58 +74,81 @@ def _bind_ffn(self):
             self._ffn = partial(LlamaTransformerLayerInfer._ffn, self)
             self._tpsp_ffn = self._tpsp_ffn_tp
 
-    def _bind_attention(self):
-        if "triton_fp8kv" in self.mode:
-            self._copy_kv_to_mem_cache = partial(Deepseek2TransformerLayerInfer._copy_kv_to_mem_cache_fp8, self)
-            self._token_attention_kernel = partial(
-                Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashdecoding_fp8, self
-            )
-        else:
-            self._copy_kv_to_mem_cache = partial(Deepseek2TransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-            if get_env_start_args().enable_fa3:
-                self._token_attention_kernel = partial(
-                    Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashattention, self
-                )
-            elif get_env_start_args().enable_flashinfer_decode:
-                self._token_attention_kernel = partial(
-                    Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashinfer, self
-                )
-            else:
-                self._token_attention_kernel = partial(
-                    Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashdecoding, self
-                )
-        if self.enable_cc_method:
-            if "triton_fp8kv" in self.mode:
-                if get_env_start_args().enable_flashinfer_prefill:
-                    self._context_attention_kernel = partial(
-                        Deepseek2TransformerLayerInfer._context_attention_flashinfer_kernel_with_CC_fp8, self
-                    )
-                else:
-                    self._context_attention_kernel = partial(
-                        Deepseek2TransformerLayerInfer._context_attention_kernel_with_CC_fp8, self
-                    )
-            else:
-                if get_env_start_args().enable_fa3:
-                    self._context_attention_kernel = partial(
-                        Deepseek2TransformerLayerInfer._context_attention_flashattention_kernel_with_CC, self
-                    )
-                elif get_env_start_args().enable_flashinfer_prefill:
-                    self._context_attention_kernel = partial(
-                        Deepseek2TransformerLayerInfer._context_attention_flashinfer_kernel_with_CC, self
-                    )
-                else:
-                    self._context_attention_kernel = partial(
-                        Deepseek2TransformerLayerInfer._context_attention_kernel_with_CC, self
-                    )
-        else:
-            if "triton_fp8kv" in self.mode:
-                self._context_attention_kernel = partial(
-                    Deepseek2TransformerLayerInfer._context_attention_kernel_origin_fp8, self
-                )
-            else:
-                self._context_attention_kernel = partial(
-                    Deepseek2TransformerLayerInfer._context_attention_kernel_origin, self
-                )
+    def _context_attention_kernel(
+        self,
+        q: torch.Tensor,
+        kv,
+        infer_state: Deepseek2InferStateInfo,
+        layer_weight: Deepseek2TransformerLayerWeight,
+        out=None,
+    ) -> torch.Tensor:
+        k_nope, k_rope, v = self._decompress_kv(
+            infer_state=infer_state,
+            layer_weight=layer_weight,
+        )
+
+        o_tensor = infer_state.prefill_att_state.prefill_att(
+            q=q,
+            k=(k_nope, k_rope),
+            v=v,
+            att_control=AttControl(mla_prefill=True, mla_prefill_dict={"softmax_scale": self.softmax_scale}),
+            alloc_func=self.alloc_tensor,
+        )
+        return o_tensor
+
+    def _token_attention_kernel(
+        self,
+        q: torch.Tensor,
+        infer_state: Deepseek2InferStateInfo,
+        layer_weight: Deepseek2TransformerLayerWeight,
+        out=None,
+    ):
+        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
+        q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
+        kv = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_)
+
+        out = infer_state.decode_att_state.decode_att(
+            q=(q_nope, q_rope),
+            k=kv,
+            v=None,
+            att_control=AttControl(mla_decode=True, mla_decode_dict={"softmax_scale": self.softmax_scale}),
+            alloc_func=self.alloc_tensor,
+        )
+        return out
+
+    def _decompress_kv(
+        self,
+        infer_state: Deepseek2InferStateInfo,
+        layer_weight: Deepseek2TransformerLayerWeight,
+    ):
+        compressed_kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
+
+        total_token_num = infer_state.total_token_num
+        sampled_compressed_kv_nope = self.alloc_tensor(
+            [total_token_num, 1, layer_weight.kv_lora_rank], dtype=compressed_kv.dtype
+        )
+        sampled_k_rope = self.alloc_tensor([total_token_num, 1, self.qk_rope_head_dim], dtype=compressed_kv.dtype)
+        sample_kv(
+            all_compressed_kv=compressed_kv,
+            sampled_compressed_kv_nope=sampled_compressed_kv_nope,
+            sampled_k_rope=sampled_k_rope,
+            b_req_idx=infer_state.b_req_idx,
+            req_to_token_indexs=infer_state.req_manager.req_to_token_indexs,
+            b_seq_len=infer_state.b_seq_len,
+            b_kv_start_loc=infer_state.b1_cu_kv_seq_len[:-1],
+            max_kv_seq_len=infer_state.max_kv_seq_len,
+        )
+        # CC
+        sampled_compressed_kv_nope = sampled_compressed_kv_nope.view(
+            total_token_num, layer_weight.kv_lora_rank
+        ).contiguous()
+        sampled_kv_nope = self.alloc_tensor(
+            [total_token_num, self.tp_q_head_num_, (self.qk_nope_head_dim + self.v_head_dim)],
+            dtype=sampled_compressed_kv_nope.dtype,
+        )
+        layer_weight.cc_kv_b_proj_.mm(sampled_compressed_kv_nope, out=sampled_kv_nope.view(total_token_num, -1))
+        sampled_k_nope, sampled_v = torch.split(sampled_kv_nope, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+        return sampled_k_nope, sampled_k_rope, sampled_v
 
     def _get_qkv(
         self,
@@ -157,20 +165,14 @@ def _get_qkv(
             q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split(
                 [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
             )
-            q = layer_weight.q_a_layernorm_.rmsnorm_forward(
-                input=q,
-                eps=self.eps_,
-                alloc_func=self.alloc_tensor,
-            )
+            q = layer_weight.q_a_layernorm_(input=q, eps=self.eps_, alloc_func=self.alloc_tensor)
             q = layer_weight.q_b_proj_.mm(q)
             cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank + self.qk_rope_head_dim)
         q = q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim)
         q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
 
-        layer_weight.kv_a_layernorm_.rmsnorm_forward(
-            cache_kv[:, :, : self.kv_lora_rank],
-            eps=self.eps_,
-            out=cache_kv[:, :, : self.kv_lora_rank],
+        layer_weight.kv_a_layernorm_(
+            cache_kv[:, :, : self.kv_lora_rank], eps=self.eps_, out=cache_kv[:, :, : self.kv_lora_rank]
         )
 
         rotary_emb_fwd(
@@ -200,10 +202,8 @@ def _tpsp_get_qkv(
             cache_kv = layer_weight.kv_a_proj_with_mqa_.mm(input).view(-1, 1, self.kv_lora_rank + self.qk_rope_head_dim)
             q = q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim)
             q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
-            layer_weight.kv_a_layernorm_.rmsnorm_forward(
-                cache_kv[:, :, : self.kv_lora_rank],
-                eps=self.eps_,
-                out=cache_kv[:, :, : self.kv_lora_rank],
+            layer_weight.kv_a_layernorm_(
+                cache_kv[:, :, : self.kv_lora_rank], eps=self.eps_, out=cache_kv[:, :, : self.kv_lora_rank]
             )
             rotary_emb_fwd(
                 q_rope,
@@ -236,19 +236,13 @@ def _tpsp_get_qkv(
                 position_sin = infer_state.position_sin
 
             q, cache_kv = qkv.split([self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1)
-            q = layer_weight.q_a_layernorm_.rmsnorm_forward(
-                q,
-                eps=self.eps_,
-                alloc_func=self.alloc_tensor,
-            )
+            q = layer_weight.q_a_layernorm_(input=q, eps=self.eps_, alloc_func=self.alloc_tensor)
             q = layer_weight.q_b_proj_.mm(q)
             cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank + self.qk_rope_head_dim)
             q = q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim)
             q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
-            layer_weight.kv_a_layernorm_.rmsnorm_forward(
-                cache_kv[:, :, : self.kv_lora_rank],
-                eps=self.eps_,
-                out=cache_kv[:, :, : self.kv_lora_rank],
+            layer_weight.kv_a_layernorm_(
+                cache_kv[:, :, : self.kv_lora_rank], eps=self.eps_, out=cache_kv[:, :, : self.kv_lora_rank]
             )
             rotary_emb_fwd(
                 q_rope,
@@ -297,423 +291,6 @@ def _tpsp_get_o(
 
         return o_tensor
 
-    def _decompress_kv(
-        self,
-        kv,
-        infer_state: Deepseek2InferStateInfo,
-        layer_weight: Deepseek2TransformerLayerWeight,
-        is_fp8,
-        total_token_num,
-        b_seq_len,
-        max_seq_len,
-        b_kv_start_loc,
-        skip_sample=False,
-    ):
-        if not skip_sample:
-            if is_fp8:
-                kv = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, :-2].view(torch.float8_e4m3fn)
-                kv_scale = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, -2:].view(torch.bfloat16)
-                k_scale = self.alloc_tensor([total_token_num, 1], dtype=kv_scale.dtype)
-            else:
-                kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-                kv_scale = None
-                k_scale = None
-
-            compressed_kv = self.alloc_tensor([total_token_num, 1, layer_weight.kv_lora_rank], dtype=kv.dtype)
-            k_rope = self.alloc_tensor([total_token_num, 1, self.qk_rope_head_dim], dtype=kv.dtype)
-            sample_kv(
-                kv,
-                compressed_kv,
-                k_rope,
-                infer_state.b_req_idx,
-                max_seq_len,
-                b_seq_len,
-                infer_state.req_manager.req_to_token_indexs,
-                b_kv_start_loc,
-                kv_scale,
-                k_scale,
-            )
-            if k_scale is not None:
-                compressed_kv = compressed_kv.to(k_scale.dtype) * k_scale.unsqueeze(-1)
-                k_rope = k_rope.to(k_scale.dtype) * k_scale.unsqueeze(-1)
-        else:
-            compressed_kv, k_rope = torch.split(  # (b*s, 1, kv_lora + qk_r)
-                kv, [layer_weight.kv_lora_rank, layer_weight.qk_rope_head_dim], dim=-1
-            )
-
-        # CC
-        compressed_kv = compressed_kv.view(-1, layer_weight.kv_lora_rank).contiguous()
-        kv_nope = self.alloc_tensor(
-            [compressed_kv.shape[0], self.tp_q_head_num_, (self.qk_nope_head_dim + self.v_head_dim)],
-            dtype=compressed_kv.dtype,
-        )
-        layer_weight.cc_kv_b_proj_.mm(compressed_kv, out=kv_nope.reshape(compressed_kv.shape[0], -1))
-        k_nope, v = torch.split(kv_nope, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
-        return k_nope, k_rope, v
-
-    # Adapted from:
-    # https://github.com/sgl-project/sglang/blob/c998d04b46920f06d945fbef9023884a768723fc/python/sglang/srt/models/deepseek_v2.py#L962
-    def _context_attention_flashattention_kernel_with_CC(
-        self,
-        q: torch.Tensor,
-        kv,
-        infer_state: Deepseek2FlashAttentionStateInfo,
-        layer_weight: Deepseek2TransformerLayerWeight,
-        out=None,
-    ) -> torch.Tensor:
-        k_nope, k_rope, v = self._decompress_kv(
-            kv,
-            infer_state,
-            layer_weight,
-            False,
-            infer_state.total_token_num,
-            infer_state.b_seq_len,
-            infer_state.max_value_in_b_seq_len,
-            infer_state.b1_kv_start_loc,
-            skip_sample=True,
-        )
-        k = torch.cat([k_nope, torch.repeat_interleave(k_rope, self.tp_q_head_num_, dim=-2)], dim=-1)
-        o_tensor, lse, *rest = flash_attn_varlen_func(
-            q=q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim),
-            k=k.view(-1, self.tp_k_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim),
-            v=v.view(-1, self.tp_v_head_num_, self.v_head_dim),
-            cu_seqlens_q=infer_state.cu_seqlens_q,
-            cu_seqlens_k=infer_state.cu_seqlens_q,
-            max_seqlen_q=infer_state.q_max_seq_len,
-            max_seqlen_k=infer_state.max_seq_len,
-            softmax_scale=self.softmax_scale,
-            causal=True,
-            return_softmax_lse=True,
-        )
-        if infer_state.has_prefix_kv:
-            k_nope, k_rope, v = self._decompress_kv(
-                kv,
-                infer_state,
-                layer_weight,
-                False,
-                infer_state.prefix_total_token_num,
-                infer_state.b_ready_cache_len,
-                infer_state.prefix_k_max_len,
-                infer_state.cu_seqlens_prefix_k,
-            )
-            k = torch.cat([k_nope, torch.repeat_interleave(k_rope, self.tp_q_head_num_, dim=-2)], dim=-1)
-            prefix_output, prefix_lse, *rest = flash_attn_varlen_func(
-                q=q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim),
-                k=k.view(-1, self.tp_k_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim),
-                v=v.view(-1, self.tp_v_head_num_, self.v_head_dim),
-                cu_seqlens_q=infer_state.cu_seqlens_q,
-                cu_seqlens_k=infer_state.cu_seqlens_prefix_k,
-                max_seqlen_q=infer_state.q_max_seq_len,
-                max_seqlen_k=infer_state.prefix_k_max_len,
-                softmax_scale=self.softmax_scale,
-                causal=False,
-                return_softmax_lse=True,
-            )
-            lse = torch.transpose(lse, 0, 1).contiguous()
-            prefix_lse = torch.transpose(prefix_lse, 0, 1).contiguous()
-            tmp_output = (
-                self.alloc_tensor((q.shape[0], q.shape[1], self.qk_nope_head_dim), dtype=q.dtype)
-                if out is None
-                else out
-            )
-            tmp_lse = torch.empty_like(lse)
-            merge_state_v2(prefix_output, prefix_lse, o_tensor, lse, tmp_output, tmp_lse)
-            o_tensor = tmp_output
-        return o_tensor
-
-    def _context_attention_flashinfer_kernel_with_CC(
-        self,
-        q: torch.Tensor,
-        kv,
-        infer_state: Deepseek2FlashInferStateInfo,
-        layer_weight: Deepseek2TransformerLayerWeight,
-        out=None,
-    ) -> torch.Tensor:
-        k_nope, k_rope, v = self._decompress_kv(
-            kv,
-            infer_state,
-            layer_weight,
-            False,
-            infer_state.total_token_num,
-            infer_state.b_seq_len,
-            infer_state.max_value_in_b_seq_len,
-            infer_state.b1_kv_start_loc,
-        )
-        o_tensor = (
-            self.alloc_tensor((q.shape[0], q.shape[1], self.qk_nope_head_dim), dtype=q.dtype) if out is None else out
-        )
-        k = torch.cat([k_nope, torch.repeat_interleave(k_rope, self.tp_q_head_num_, dim=-2)], dim=-1)
-        infer_state.prefill_wrapper.run(q, k, v, out=o_tensor)
-        return o_tensor
-
-    def _context_attention_flashinfer_kernel_with_CC_fp8(
-        self,
-        q: torch.Tensor,
-        kv,
-        infer_state: Deepseek2FlashInferStateInfo,
-        layer_weight: Deepseek2TransformerLayerWeight,
-        out=None,
-    ) -> torch.Tensor:
-        k_nope, k_rope, v = self._decompress_kv(
-            kv,
-            infer_state,
-            layer_weight,
-            True,
-            infer_state.total_token_num,
-            infer_state.b_seq_len,
-            infer_state.max_value_in_b_seq_len,
-            infer_state.b1_kv_start_loc,
-        )
-        o_tensor = (
-            self.alloc_tensor((q.shape[0], q.shape[1], self.qk_nope_head_dim), dtype=q.dtype) if out is None else out
-        )
-        k = torch.cat([k_nope, torch.repeat_interleave(k_rope, self.tp_q_head_num_, dim=-2)], dim=-1)
-        infer_state.prefill_wrapper.run(q, k, v, out=o_tensor)
-        return o_tensor
-
-    def _context_attention_kernel_with_CC(
-        self,
-        q: torch.Tensor,
-        kv,
-        infer_state: Deepseek2InferStateInfo,
-        layer_weight: Deepseek2TransformerLayerWeight,
-        out=None,
-    ) -> torch.Tensor:
-        k_nope, k_rope, v = self._decompress_kv(
-            kv,
-            infer_state,
-            layer_weight,
-            False,
-            infer_state.total_token_num,
-            infer_state.b_seq_len,
-            infer_state.max_value_in_b_seq_len,
-            infer_state.b1_kv_start_loc,
-        )
-        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
-        o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype) if out is None else out
-        context_attention_fwd_with_v(
-            q_nope,
-            q_rope,
-            k_nope,
-            k_rope,
-            v,
-            o_tensor.view(-1, self.tp_q_head_num_, q_nope.shape[-1]),
-            infer_state.b_start_loc,
-            infer_state.b1_kv_start_loc,
-            infer_state.b_seq_len,
-            infer_state.b_ready_cache_len,
-            infer_state.max_len_in_batch,
-            self.softmax_scale,
-        )
-        return o_tensor
-
-    def _context_attention_kernel_with_CC_fp8(
-        self,
-        q: torch.Tensor,
-        kv,
-        infer_state: Deepseek2InferStateInfo,
-        layer_weight: Deepseek2TransformerLayerWeight,
-        out=None,
-    ) -> torch.Tensor:
-        k_nope, k_rope, v = self._decompress_kv(
-            kv,
-            infer_state,
-            layer_weight,
-            True,
-            infer_state.total_token_num,
-            infer_state.b_seq_len,
-            infer_state.max_value_in_b_seq_len,
-            infer_state.b1_kv_start_loc,
-        )
-        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
-        o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype) if out is None else out
-        context_attention_fwd_with_v(
-            q_nope,
-            q_rope,
-            k_nope,
-            k_rope,
-            v,
-            o_tensor.view(-1, self.tp_q_head_num_, q_nope.shape[-1]),
-            infer_state.b_start_loc,
-            infer_state.b1_kv_start_loc,
-            infer_state.b_seq_len,
-            infer_state.b_ready_cache_len,
-            infer_state.max_len_in_batch,
-            self.softmax_scale,
-        )
-        return o_tensor
-
-    def _context_attention_kernel_origin(
-        self,
-        q: torch.Tensor,
-        kv,
-        infer_state: Deepseek2InferStateInfo,
-        layer_weight: Deepseek2TransformerLayerWeight,
-        out=None,
-    ) -> torch.Tensor:
-        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
-        q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
-        o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype) if out is None else out
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        context_attention_fwd(
-            q_nope,
-            q_rope,
-            kv[:, :, : -self.qk_rope_head_dim],
-            kv[:, :, -self.qk_rope_head_dim :],
-            o_tensor.view(-1, self.tp_q_head_num_, self.kv_lora_rank),
-            infer_state.b_req_idx,
-            infer_state.b_start_loc,
-            infer_state.b_seq_len,
-            infer_state.b_ready_cache_len,
-            infer_state.max_len_in_batch,
-            infer_state.req_manager.req_to_token_indexs,
-            self.softmax_scale,
-        )
-        return o_tensor
-
-    def _context_attention_kernel_origin_fp8(
-        self,
-        q: torch.Tensor,
-        kv,
-        infer_state: Deepseek2InferStateInfo,
-        layer_weight: Deepseek2TransformerLayerWeight,
-        out=None,
-    ) -> torch.Tensor:
-        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
-        q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
-        o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype) if out is None else out
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, :-2].view(torch.float8_e4m3fn)
-        kv_scale = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, -2:].view(torch.bfloat16)
-        context_attention_fwd_fp8(
-            q_nope,
-            q_rope,
-            kv[:, :, : -self.qk_rope_head_dim],
-            kv[:, :, -self.qk_rope_head_dim :],
-            kv_scale,
-            o_tensor.view(-1, self.tp_q_head_num_, self.kv_lora_rank),
-            infer_state.b_req_idx,
-            infer_state.b_start_loc,
-            infer_state.b_seq_len,
-            infer_state.b_ready_cache_len,
-            infer_state.max_len_in_batch,
-            infer_state.req_manager.req_to_token_indexs,
-            self.softmax_scale,
-        )
-        return o_tensor
-
-    def _token_gqa_decode_attention_flashattention(
-        self, q, infer_state: Deepseek2FlashAttentionStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
-    ):
-        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
-        q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        k_rope = kv[:, :, -self.qk_rope_head_dim :].reshape(-1, 1, 1, self.qk_rope_head_dim)
-        kv_nope = kv[:, :, : -self.qk_rope_head_dim].reshape(-1, 1, 1, self.kv_lora_rank)
-        k_descale, v_descale = None, None
-        o_tensor = flash_attn_with_kvcache(
-            q=q_rope,
-            k_cache=k_rope,
-            v_cache=kv_nope,
-            qv=q_nope,
-            page_table=infer_state.page_table,
-            cache_seqlens=infer_state.b_att_seq_len,
-            cu_seqlens_q=infer_state.cu_seqlens_q,
-            cu_seqlens_k_new=infer_state.cu_seqlens_k,
-            max_seqlen_q=infer_state.max_q_seq_len,
-            softmax_scale=self.softmax_scale,
-            causal=True,
-            window_size=(-1, -1),
-            softcap=0.0,
-            k_descale=k_descale,
-            v_descale=v_descale,
-            return_softmax_lse=False,
-        )
-        return o_tensor
-
-    def _token_gqa_decode_attention_flashinfer(
-        self, q, infer_state: Deepseek2FlashInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
-    ):
-        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
-        q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
-
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype)
-
-        infer_state.decode_wrapper.run(
-            q_nope,
-            q_rope,
-            kv[:, :, : -self.qk_rope_head_dim],
-            kv[:, :, -self.qk_rope_head_dim :],
-            out=o_tensor,
-            return_lse=False,
-        )
-        return o_tensor
-
-    def _token_gqa_decode_attention_flashdecoding(
-        self, q, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
-    ):
-        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
-        q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        out = gqa_token_decode_attention_flash_decoding(
-            q_nope,
-            q_rope,
-            kv[:, :, : -self.qk_rope_head_dim],
-            kv[:, :, -self.qk_rope_head_dim :],
-            infer_state,
-            self.tp_q_head_num_,
-            self.kv_lora_rank,
-            self.qk_rope_head_dim,
-            self.qk_nope_head_dim,
-            self.softmax_scale,
-            alloc_tensor_func=self.alloc_tensor,
-        )
-        return out
-
-    def _token_gqa_decode_attention_flashdecoding_fp8(
-        self, q, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
-    ):
-        q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
-        q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
-
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, :-2].view(torch.float8_e4m3fn)
-        kv_scale = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, -2:].view(torch.bfloat16)
-        return gqa_token_decode_attention_flash_decoding_fp8(
-            q_nope,
-            q_rope,
-            kv[:, :, : -self.qk_rope_head_dim],
-            kv[:, :, -self.qk_rope_head_dim :],
-            kv_scale,
-            infer_state,
-            self.tp_q_head_num_,
-            self.kv_lora_rank,
-            self.qk_rope_head_dim,
-            self.qk_nope_head_dim,
-            self.softmax_scale,
-            alloc_tensor_func=self.alloc_tensor,
-        )
-
-    def _copy_kv_to_mem_cache_normal(self, buffer, mem_index, mem_manager):
-        destindex_copy_kv(
-            buffer[:, :, : self.kv_lora_rank],
-            buffer[:, :, self.kv_lora_rank :],
-            mem_index,
-            mem_manager.kv_buffer[self.layer_num_][:, :, : self.kv_lora_rank],
-            mem_manager.kv_buffer[self.layer_num_][:, :, self.kv_lora_rank :],
-        )
-        return
-
-    def _copy_kv_to_mem_cache_fp8(self, buffer, mem_index, mem_manager):
-        destindex_copy_kv_fp8(
-            buffer[:, :, : self.kv_lora_rank],
-            buffer[:, :, self.kv_lora_rank :],
-            mem_index,
-            mem_manager.kv_buffer[self.layer_num_][:, :, : self.kv_lora_rank].view(torch.float8_e4m3fn),
-            mem_manager.kv_buffer[self.layer_num_][:, :, self.kv_lora_rank : -2].view(torch.float8_e4m3fn),
-            mem_manager.kv_buffer[self.layer_num_][:, :, -2:].view(buffer.dtype),
-        )
-        return
-
     def _moe_ffn(
         self, input, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
     ) -> torch.Tensor:
diff --git a/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py b/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py
index 611878f9e8..783e70e64b 100644
--- a/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py
@@ -6,20 +6,18 @@
 from lightllm.utils.envs_utils import enable_env_vars, get_env_start_args
 from lightllm.common.basemodel.layer_weights.meta_weights import (
     ROWMMWeight,
-    COLMMWeight,
-    NoTpNormWeight,
-    FusedMoeWeightEP,
     ROWBMMWeight,
-    create_tp_moe_wegiht_obj,
+    COLMMWeight,
+    RMSNormWeight,
+    FusedMoeWeight,
 )
-from functools import partial
 from ..triton_kernel.weight_dequant import weight_dequant
 
 
 class Deepseek2TransformerLayerWeight(TransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
         self.enable_cc_method = not os.getenv("DISABLE_CC_METHOD", "False").upper() in ["ON", "TRUE", "1"]
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _parse_config(self):
@@ -40,10 +38,16 @@ def _parse_config(self):
         self.kv_lora_rank = self.network_config_["kv_lora_rank"]
         self.num_fused_shared_experts = 0
         if get_env_start_args().enable_fused_shared_experts and self.is_moe:
-            # MOE_MODE 处于 TP 模式下才能使能 enable_fused_shared_experts
-            moe_mode = os.getenv("MOE_MODE", "TP")
-            assert moe_mode == "TP"
+            # enable_fused_shared_experts can only work with tensor parallelism
+            assert not get_env_start_args().enable_ep_moe, "enable_fused_shared_experts can only work with tp mode."
             self.num_fused_shared_experts = self.network_config_.get("n_shared_experts", 0)
+        self.n_embed = self.network_config_["hidden_size"]
+        self.n_inter = self.network_config_["intermediate_size"]
+        self.moe_inter = self.network_config_.get("moe_intermediate_size", self.n_inter)
+        self.q_out_dim = self.num_attention_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim)
+        self.kv_a_out_dim = self.kv_lora_rank + self.qk_rope_head_dim
+        self.kv_b_out_dim = self.num_attention_heads * (self.qk_nope_head_dim + self.v_head_dim)
+        self.o_in_dim = self.num_attention_heads * self.v_head_dim
 
     def _init_weight_names(self):
         if self.q_lora_rank is None:
@@ -60,31 +64,14 @@ def _init_weight(self):
             self._init_ffn()
         self._init_norm()
 
-    def _load_kb(self, kv_b_proj_):
-        k_b_proj_ = kv_b_proj_.view(self.num_attention_heads, self.qk_nope_head_dim * 2, self.kv_lora_rank)[
-            :, : self.qk_nope_head_dim, :
-        ]
-        return k_b_proj_.contiguous().to(kv_b_proj_.dtype)
-
-    def _load_kb_scale(self, kv_b_proj_, block_size):
-        k_b_proj_scale_ = kv_b_proj_.view(
-            self.num_attention_heads, self.qk_nope_head_dim * 2 // block_size, self.kv_lora_rank // block_size
-        )[:, : self.qk_nope_head_dim // block_size, :]
-        return k_b_proj_scale_.contiguous().to(kv_b_proj_.dtype)
-
-    def _load_vb(self, kv_b_proj_):
-        v_b_proj_ = kv_b_proj_.T.view(self.kv_lora_rank, self.num_attention_heads, self.qk_nope_head_dim * 2,)[
-            :, :, self.qk_nope_head_dim :
-        ].transpose(0, 1)
-        return v_b_proj_.contiguous().to(kv_b_proj_.dtype)
-
-    def _load_vb_scale(self, kv_b_proj_scale_, block_size):
-        v_b_proj_scale_ = kv_b_proj_scale_.T.view(
-            self.kv_lora_rank // block_size,
-            self.num_attention_heads,
-            self.qk_nope_head_dim * 2 // block_size,
-        )[:, :, self.qk_nope_head_dim // block_size :].transpose(0, 1)
-        return v_b_proj_scale_.contiguous().to(kv_b_proj_scale_.dtype)
+    def _split_kv_b_proj(self, kv_b_proj_):
+        kv_b_proj_ = kv_b_proj_.view(self.num_attention_heads, self.qk_nope_head_dim * 2, self.kv_lora_rank)
+        k_b_proj_, v_b_proj_ = torch.split(kv_b_proj_, [self.qk_nope_head_dim, self.v_head_dim], dim=-2)
+        # num_attention_heads x qk_nope_head_dim x kv_lora_rank
+        k_b_proj_ = k_b_proj_.contiguous().to(kv_b_proj_.dtype)
+        # num_attention_heads x kv_lora_rank x v_head_dim
+        v_b_proj_ = v_b_proj_.transpose(1, 2).contiguous().to(kv_b_proj_.dtype)
+        return k_b_proj_, v_b_proj_
 
     def _rename_shared_experts(self, weights, weight_scale_suffix):
         # 将共享专家对应的参数，改造为与路由专家一致的权重名称和映射关系。
@@ -108,7 +95,6 @@ def load_hf_weights(self, weights):
         weight_scale_suffix = None
         if self.quant_cfg.quantized_weight:
             weight_scale_suffix = kv_b_quant_method.weight_scale_suffix
-
         if f"model.layers.{self.layer_num_}.self_attn.kv_b_proj.weight" in weights:
             kv_b_proj_ = weights[f"model.layers.{self.layer_num_}.self_attn.kv_b_proj.weight"]
             # for deepseek_v3, the bmm operator is not quantized
@@ -117,21 +103,9 @@ def load_hf_weights(self, weights):
                     kv_b_proj_.cuda(),
                     weights[f"model.layers.{self.layer_num_}.self_attn.kv_b_proj." + weight_scale_suffix].cuda(),
                 ).cpu()
-            weights[f"model.layers.{self.layer_num_}.self_attn.k_b_proj.weight"] = self._load_kb(kv_b_proj_)
-            weights[f"model.layers.{self.layer_num_}.self_attn.v_b_proj.weight"] = self._load_vb(kv_b_proj_)
-
-        if (
-            self.quant_cfg.quantized_weight
-            and f"model.layers.{self.layer_num_}.self_attn.kv_b_proj." + weight_scale_suffix in weights
-        ):
-            kv_b_proj_scale_ = weights[f"model.layers.{self.layer_num_}.self_attn.kv_b_proj." + weight_scale_suffix]
-            block_size = 128
-            weights[f"model.layers.{self.layer_num_}.self_attn.k_b_proj." + weight_scale_suffix] = self._load_kb_scale(
-                kv_b_proj_scale_, block_size
-            )
-            weights[f"model.layers.{self.layer_num_}.self_attn.v_b_proj." + weight_scale_suffix] = self._load_vb_scale(
-                kv_b_proj_scale_, block_size
-            )
+            k_b_proj_, v_b_proj_ = self._split_kv_b_proj(kv_b_proj_)
+            weights[f"model.layers.{self.layer_num_}.self_attn.k_b_proj.weight"] = k_b_proj_
+            weights[f"model.layers.{self.layer_num_}.self_attn.v_b_proj.weight"] = v_b_proj_
 
         # rename the shared experts weight
         if self.num_fused_shared_experts > 0:
@@ -141,116 +115,120 @@ def load_hf_weights(self, weights):
     def _init_qkvo(self):
         if self.q_lora_rank is None:
             self.q_weight_ = ROWMMWeight(
+                in_dim=self.n_embed,
+                out_dims=[self.q_out_dim],
                 weight_names=f"model.layers.{self.layer_num_}.self_attn.q_proj.weight",
                 data_type=self.data_type_,
-                quant_cfg=self.quant_cfg,
-                layer_num=self.layer_num_,
-                name="q_weight",
+                quant_method=self.get_quant_method("q_weight"),
             )
             self.kv_a_proj_with_mqa_ = ROWMMWeight(
+                in_dim=self.n_embed,
+                out_dims=[self.kv_a_out_dim],
                 weight_names=f"model.layers.{self.layer_num_}.self_attn.kv_a_proj_with_mqa.weight",
                 data_type=self.data_type_,
-                quant_cfg=self.quant_cfg,
-                layer_num=self.layer_num_,
-                name="kv_a_proj_with_mqa",
+                quant_method=self.get_quant_method("kv_a_proj_with_mqa"),
                 tp_rank=0,
                 tp_world_size=1,
             )
         else:
             self.qkv_a_proj_with_mqa_ = ROWMMWeight(
+                in_dim=self.n_embed,
+                out_dims=[self.q_lora_rank, self.kv_a_out_dim],
                 weight_names=[
                     f"model.layers.{self.layer_num_}.self_attn.q_a_proj.weight",
                     f"model.layers.{self.layer_num_}.self_attn.kv_a_proj_with_mqa.weight",
                 ],
                 data_type=self.data_type_,
-                quant_cfg=self.quant_cfg,
-                layer_num=self.layer_num_,
-                name="qkv_a_proj_with_mqa",
+                quant_method=self.get_quant_method("qkv_a_proj_with_mqa"),
                 tp_rank=0,
                 tp_world_size=1,
             )
             self.q_b_proj_ = ROWMMWeight(
+                in_dim=self.q_lora_rank,
+                out_dims=[self.q_out_dim],
                 weight_names=f"model.layers.{self.layer_num_}.self_attn.q_b_proj.weight",
                 data_type=self.data_type_,
-                quant_cfg=self.quant_cfg,
-                layer_num=self.layer_num_,
-                name="q_b_proj",
+                quant_method=self.get_quant_method("q_b_proj"),
             )
         self.k_b_proj_ = ROWBMMWeight(
+            dim0=self.num_attention_heads,
+            dim1=self.qk_nope_head_dim,
+            dim2=self.kv_lora_rank,
             weight_names=f"model.layers.{self.layer_num_}.self_attn.k_b_proj.weight",
             data_type=self.data_type_,
-            quant_cfg=None,
-            layer_num=self.layer_num_,
-            name="k_b_proj",
+            quant_method=None,
         )
         self.v_b_proj_ = ROWBMMWeight(
+            dim0=self.num_attention_heads,
+            dim1=self.kv_lora_rank,
+            dim2=self.v_head_dim,
             weight_names=f"model.layers.{self.layer_num_}.self_attn.v_b_proj.weight",
             data_type=self.data_type_,
-            quant_cfg=None,
-            layer_num=self.layer_num_,
-            name="v_b_proj",
+            quant_method=None,
         )
         if self.enable_cc_method:
             self.cc_kv_b_proj_ = ROWMMWeight(
+                in_dim=self.kv_lora_rank,
+                out_dims=[self.kv_b_out_dim],
                 weight_names=f"model.layers.{self.layer_num_}.self_attn.kv_b_proj.weight",
                 data_type=self.data_type_,
-                quant_cfg=self.quant_cfg,
-                layer_num=self.layer_num_,
-                name="cc_kv_b_proj",
+                quant_method=self.get_quant_method("cc_kv_b_proj"),
             )
 
         self.o_weight_ = COLMMWeight(
+            in_dim=self.o_in_dim,
+            out_dims=[self.n_embed],
             weight_names=f"model.layers.{self.layer_num_}.self_attn.o_proj.weight",
             data_type=self.data_type_,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="o_weight",
+            quant_method=self.get_quant_method("o_weight"),
         )
 
-    def _load_mlp(self, mlp_prefix):
-        moe_mode = os.getenv("MOE_MODE", "TP")
-        if self.is_moe and moe_mode == "EP":
+    def _load_mlp(self, mlp_prefix, is_shared_experts=False):
+        enable_ep_moe = get_env_start_args().enable_ep_moe
+        mlp_inter = self.moe_inter if is_shared_experts else self.n_inter
+        if self.is_moe and enable_ep_moe:
             self.gate_up_proj = ROWMMWeight(
+                in_dim=self.n_embed,
+                out_dims=[mlp_inter, mlp_inter],
                 weight_names=[f"{mlp_prefix}.gate_proj.weight", f"{mlp_prefix}.up_proj.weight"],
                 data_type=self.data_type_,
-                quant_cfg=self.quant_cfg,
-                layer_num=self.layer_num_,
-                name="gate_up_proj",
+                quant_method=self.get_quant_method("gate_up_proj"),
                 tp_rank=0,
                 tp_world_size=1,
             )
             self.down_proj = COLMMWeight(
+                in_dim=mlp_inter,
+                out_dims=[self.n_embed],
                 weight_names=f"{mlp_prefix}.down_proj.weight",
                 data_type=self.data_type_,
-                quant_cfg=self.quant_cfg,
-                layer_num=self.layer_num_,
-                name="down_proj",
+                quant_method=self.get_quant_method("down_proj"),
                 tp_rank=0,
                 tp_world_size=1,
             )
         else:
             self.gate_up_proj = ROWMMWeight(
+                in_dim=self.n_embed,
+                out_dims=[mlp_inter, mlp_inter],
                 weight_names=[f"{mlp_prefix}.gate_proj.weight", f"{mlp_prefix}.up_proj.weight"],
                 data_type=self.data_type_,
-                quant_cfg=self.quant_cfg,
-                layer_num=self.layer_num_,
-                name="gate_up_proj",
+                quant_method=self.get_quant_method("gate_up_proj"),
             )
             self.down_proj = COLMMWeight(
+                in_dim=mlp_inter,
+                out_dims=[self.n_embed],
                 weight_names=f"{mlp_prefix}.down_proj.weight",
                 data_type=self.data_type_,
-                quant_cfg=self.quant_cfg,
-                layer_num=self.layer_num_,
-                name="down_proj",
+                quant_method=self.get_quant_method("down_proj"),
             )
 
     def _init_moe(self):
         moe_intermediate_size = self.network_config_["moe_intermediate_size"]
         self.moe_gate = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_routed_experts],
             weight_names=f"model.layers.{self.layer_num_}.mlp.gate.weight",
             data_type=self.data_type_,
-            layer_num=self.layer_num_,
-            name="moe_gate",
+            quant_method=None,
             tp_rank=0,
             tp_world_size=1,
         )
@@ -261,54 +239,47 @@ def _init_moe(self):
         # 专家对应的 gate_up_proj 等weight 参数。当 num_fused_shared_experts
         # == 0 时，说明不存在融合共享专家，共享专家单独加载和进行推理。
         if self.num_fused_shared_experts == 0:
-            self._load_mlp(f"model.layers.{self.layer_num_}.mlp.shared_experts")
-        moe_mode = os.getenv("MOE_MODE", "TP")
-        assert moe_mode in ["EP", "TP"]
-        if moe_mode == "TP":
-            self.experts = create_tp_moe_wegiht_obj(
-                gate_proj_name="gate_proj",
-                down_proj_name="down_proj",
-                up_proj_name="up_proj",
-                e_score_correction_bias_name=self.e_score_correction_bias_name,
-                weight_prefix=f"model.layers.{self.layer_num_}.mlp.experts",
-                n_routed_experts=self.n_routed_experts,
-                num_fused_shared_experts=self.num_fused_shared_experts,
-                split_inter_size=moe_intermediate_size // self.tp_world_size_,
-                data_type=self.data_type_,
-                network_config=self.network_config_,
-                layer_num=self.layer_num_,
-                quant_cfg=self.quant_cfg,
-            )
-        elif moe_mode == "EP":
-            self.experts = FusedMoeWeightEP(
-                gate_proj_name="gate_proj",
-                down_proj_name="down_proj",
-                up_proj_name="up_proj",
-                e_score_correction_bias_name=self.e_score_correction_bias_name,
-                weight_prefix=f"model.layers.{self.layer_num_}.mlp.experts",
-                n_routed_experts=self.n_routed_experts,
-                data_type=self.data_type_,
-                network_config=self.network_config_,
-                layer_num=self.layer_num_,
-                quant_cfg=self.quant_cfg,
-            )
-        else:
-            raise ValueError(f"Unsupported moe mode: {moe_mode}")
+            self._load_mlp(f"model.layers.{self.layer_num_}.mlp.shared_experts", is_shared_experts=True)
+        self.experts = FusedMoeWeight(
+            gate_proj_name="gate_proj",
+            down_proj_name="down_proj",
+            up_proj_name="up_proj",
+            e_score_correction_bias_name=self.e_score_correction_bias_name,
+            weight_prefix=f"model.layers.{self.layer_num_}.mlp.experts",
+            n_routed_experts=self.n_routed_experts,
+            hidden_size=self.n_embed,
+            moe_intermediate_size=moe_intermediate_size,
+            data_type=self.data_type_,
+            quant_method=self.quant_cfg.get_quant_method(self.layer_num_, "fused_moe"),
+            num_fused_shared_experts=self.num_fused_shared_experts,
+            layer_num=self.layer_num_,
+            network_config=self.network_config_,
+        )
 
     def _init_ffn(self):
         self._load_mlp(f"model.layers.{self.layer_num_}.mlp")
 
     def _init_norm(self):
-        self.att_norm_weight_ = NoTpNormWeight(
-            f"model.layers.{self.layer_num_}.input_layernorm.weight", self.data_type_
+        hidden_size = self.network_config_["hidden_size"]
+
+        self.att_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
+            weight_name=f"model.layers.{self.layer_num_}.input_layernorm.weight",
+            data_type=self.data_type_,
         )
-        self.ffn_norm_weight_ = NoTpNormWeight(
-            f"model.layers.{self.layer_num_}.post_attention_layernorm.weight", self.data_type_
+        self.ffn_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
+            weight_name=f"model.layers.{self.layer_num_}.post_attention_layernorm.weight",
+            data_type=self.data_type_,
         )
-        self.kv_a_layernorm_ = NoTpNormWeight(
-            f"model.layers.{self.layer_num_}.self_attn.kv_a_layernorm.weight", self.data_type_
+        self.kv_a_layernorm_ = RMSNormWeight(
+            dim=self.kv_lora_rank,
+            weight_name=f"model.layers.{self.layer_num_}.self_attn.kv_a_layernorm.weight",
+            data_type=self.data_type_,
         )
         if self.q_lora_rank is not None:
-            self.q_a_layernorm_ = NoTpNormWeight(
-                f"model.layers.{self.layer_num_}.self_attn.q_a_layernorm.weight", self.data_type_
+            self.q_a_layernorm_ = RMSNormWeight(
+                dim=self.q_lora_rank,
+                weight_name=f"model.layers.{self.layer_num_}.self_attn.q_a_layernorm.weight",
+                data_type=self.data_type_,
             )
diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py
index e4ce7c8269..f0739a8a81 100644
--- a/lightllm/models/deepseek2/model.py
+++ b/lightllm/models/deepseek2/model.py
@@ -4,51 +4,16 @@
 from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer
 from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import Deepseek2TransformerLayerWeight
 from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
-from lightllm.models.deepseek2.flashinfer_struct import Deepseek2FlashInferStateInfo
-from lightllm.models.deepseek2.flashattention_infer_struct import Deepseek2FlashAttentionStateInfo
-from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights
-
 from lightllm.models.llama.model import LlamaTpPartModel
 from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class
 from lightllm.utils.log_utils import init_logger
-from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale
 from lightllm.utils.envs_utils import enable_env_vars, get_env_start_args, get_added_mtp_kv_layer_num
 from lightllm.distributed.communication_op import dist_group_manager
-from lightllm.utils.dist_utils import get_dp_world_size, get_current_device_id
-
+from lightllm.common.basemodel.attention import get_mla_decode_att_backend_class, get_mla_prefill_att_backend_class
 
 logger = init_logger(__name__)
 
 
-class DeepSeek2FlashInferStateExtraInfo:
-    def __init__(self, model):
-        num_heads = model.config["num_attention_heads"]
-        self.tp_q_head_num = num_heads // get_dp_world_size()
-        self.qk_nope_head_dim = model.qk_nope_head_dim
-        self.qk_rope_head_dim = model.qk_rope_head_dim
-        self.kv_lora_rank = model.kv_lora_rank
-        self.q_data_type = model.data_type
-        self.kv_data_type = model.data_type
-        self.workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8, device=get_current_device_id())
-        self.max_seq_length = model.max_seq_length
-        self.softmax_scale = (self.qk_nope_head_dim + self.qk_rope_head_dim) ** (-0.5)
-        self.kv_indices_buffer = [
-            torch.empty(
-                model.graph_max_batch_size * self.max_seq_length, dtype=torch.int32, device=get_current_device_id()
-            ),
-            torch.empty(
-                model.graph_max_batch_size * self.max_seq_length, dtype=torch.int32, device=get_current_device_id()
-            ),
-        ]
-        if model.config["rope_scaling"] is not None:
-            rope_scaling = model.config["rope_scaling"]
-            mscale_all_dim = rope_scaling.get("mscale_all_dim", 0)
-            scaling_factor = rope_scaling["factor"]
-            if mscale_all_dim:
-                mscale = get_deepseek_mscale(scaling_factor, mscale_all_dim)
-                self.softmax_scale = self.softmax_scale * mscale * mscale
-
-
 @ModelRegistry(["deepseek_v2", "deepseek_v3"])
 class Deepseek2TpPartModel(LlamaTpPartModel):
     # weight class
@@ -61,18 +26,13 @@ class Deepseek2TpPartModel(LlamaTpPartModel):
     infer_state_class = Deepseek2InferStateInfo
 
     def __init__(self, kvargs):
-        self.enable_flashinfer = (
-            get_env_start_args().enable_flashinfer_prefill or get_env_start_args().enable_flashinfer_decode
-        )
         super().__init__(kvargs)
         return
 
-    def _init_inferstate_cls(self):
-        if get_env_start_args().enable_fa3:
-            self.infer_state_class = Deepseek2FlashAttentionStateInfo
-        elif self.enable_flashinfer:
-            self.infer_state_class = Deepseek2FlashInferStateInfo
-            self.flashinfer_extra_state = DeepSeek2FlashInferStateExtraInfo(self)
+    def _init_att_backend(self):
+        self.prefill_att_backend = get_mla_prefill_att_backend_class(index=0)(model=self)
+        self.decode_att_backend = get_mla_decode_att_backend_class(index=0)(model=self)
+        return
 
     def _init_some_value(self):
         super()._init_some_value()
diff --git a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py b/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py
index fd437c3888..b9be73e278 100644
--- a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py
+++ b/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py
@@ -5,7 +5,6 @@
 import triton.language as tl
 from typing import List
 from lightllm.utils.log_utils import init_logger
-from .gqa_flash_decoding_config import MlaDecodeAttentionKernelConfig
 from lightllm.utils.device_utils import get_device_sm_count
 
 logger = init_logger(__name__)
@@ -28,16 +27,18 @@ def gqa_token_decode_attention_flash_decoding_fp8(
     **run_config
 ):
     batch_size = infer_state.batch_size
-    max_len_in_batch = infer_state.max_len_in_batch
+    max_kv_seq_len = infer_state.max_kv_seq_len
     calcu_shape1 = (batch_size, q_head_num, kv_lora_rank)
     calcu_shape2 = (batch_size, q_head_num, q_rope_dim)
 
     if not run_config:
         if torch.cuda.is_current_stream_capturing():
-            avg_seq_len_in_batch = max_len_in_batch
+            avg_seq_len_in_batch = max_kv_seq_len
         else:
             avg_seq_len_in_batch = infer_state.total_token_num // batch_size
 
+        from .gqa_flash_decoding_config import MlaDecodeAttentionKernelConfig
+
         run_config = MlaDecodeAttentionKernelConfig.try_to_get_best_config(
             batch_size=batch_size,
             avg_seq_len_in_batch=avg_seq_len_in_batch,
@@ -191,7 +192,7 @@ def _fwd_kernel_calcu_index_and_block_seq(
 
     infer_state = Deepseek2InferStateInfo()
     infer_state.batch_size = Z
-    infer_state.max_len_in_batch = N_CTX
+    infer_state.max_kv_seq_len = N_CTX
     infer_state.total_token_num = Z * N_CTX
     infer_state.req_manager = ReqManager(Z, N_CTX, None)
     infer_state.req_manager.req_to_token_indexs = req_to_token_indexs
diff --git a/lightllm/models/deepseek2/triton_kernel/sample_kv.py b/lightllm/models/deepseek2/triton_kernel/sample_kv.py
index af0aaa2f66..53a0a60eb2 100644
--- a/lightllm/models/deepseek2/triton_kernel/sample_kv.py
+++ b/lightllm/models/deepseek2/triton_kernel/sample_kv.py
@@ -8,111 +8,101 @@
 
 @triton.jit
 def _sample_kv_kernel(
-    KV_input,
-    KV_scale,
-    KV_nope,
-    KV_rope,
-    K_scale,
-    B_start_loc,
-    B_Seqlen,
-    Req_to_tokens,
-    B_req_idx,
-    stride_input_dim,
-    stride_scale_dim,
-    stride_nope_dim,
-    stride_rope_dim,
+    all_compressed_kv,
+    stride_all_s,
+    stride_all_d,
+    sampled_compressed_kv_nope,
+    stride_nope_s,
+    stride_nope_d,
+    sampled_k_rope,
+    stride_rope_s,
+    stride_rope_d,
+    b_kv_start_loc,
+    b_seq_len,
+    req_to_token_indexs,
     stride_req_to_tokens_b,
-    HAS_SCALE: tl.constexpr,
-    BLOCK_M: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr,
-    BLOCK_ROPE_DMODEL: tl.constexpr,
+    b_req_idx,
+    BLOCK_SEQ: tl.constexpr,
+    BLOCK_NOPE_DIM: tl.constexpr,
+    BLOCK_ROPE_DIM: tl.constexpr,
 ):
     cur_batch = tl.program_id(0)
     start_m = tl.program_id(1)
 
-    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
-    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
-    cur_batch_start_loc = tl.load(B_start_loc + cur_batch)
+    cur_batch_seq_len = tl.load(b_seq_len + cur_batch)
+    cur_batch_req_idx = tl.load(b_req_idx + cur_batch)
+    cur_batch_start_loc = tl.load(b_kv_start_loc + cur_batch)
 
-    offs_nope_d = tl.arange(0, BLOCK_DMODEL)
-    offs_rope_d = tl.arange(0, BLOCK_ROPE_DMODEL)
-    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_nope_d = tl.arange(0, BLOCK_NOPE_DIM)
+    offs_rope_d = tl.arange(0, BLOCK_ROPE_DIM)
+    offs_m = (start_m * BLOCK_SEQ + tl.arange(0, BLOCK_SEQ)) % cur_batch_seq_len
 
-    block_end_loc = tl.minimum((start_m + 1) * BLOCK_M, cur_batch_seq_len)
+    if start_m * BLOCK_SEQ > cur_batch_seq_len:
+        return
 
     kv_loc = tl.load(
-        Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_m,
-        mask=offs_m < block_end_loc,
-        other=0,
+        req_to_token_indexs + stride_req_to_tokens_b * cur_batch_req_idx + offs_m,
     ).to(tl.int64)
-    off_kv_nope = kv_loc[:, None] * stride_input_dim + offs_nope_d[None, :]
-    off_kv_rope = kv_loc[:, None] * stride_input_dim + (offs_rope_d + BLOCK_DMODEL)[None, :]
-    kv_nope = tl.load(KV_input + off_kv_nope, mask=offs_m[:, None] < block_end_loc, other=0.0)
-    kv_rope = tl.load(KV_input + off_kv_rope, mask=offs_m[:, None] < block_end_loc, other=0.0)
-    off_nope = (offs_m + cur_batch_start_loc)[:, None] * stride_nope_dim + offs_nope_d[None, :]
-    off_rope = (offs_m + cur_batch_start_loc)[:, None] * stride_rope_dim + offs_rope_d[None, :]
-    nope_ptrs = KV_nope + off_nope
-    rope_ptrs = KV_rope + off_rope
-    tl.store(nope_ptrs, kv_nope, mask=offs_m[:, None] < block_end_loc)
-    tl.store(rope_ptrs, kv_rope, mask=offs_m[:, None] < block_end_loc)
-    if HAS_SCALE:
-        kv_scale = tl.load(KV_scale + kv_loc * stride_scale_dim, mask=offs_m < block_end_loc)
-        off_k_scale = cur_batch_start_loc + offs_m
-        k_scale_ptrs = K_scale + off_k_scale
-        tl.store(k_scale_ptrs, kv_scale, mask=offs_m < block_end_loc)
+    off_kv_nope = kv_loc[:, None] * stride_all_s + offs_nope_d[None, :]
+    off_kv_rope = kv_loc[:, None] * stride_all_s + (offs_rope_d + BLOCK_NOPE_DIM)[None, :]
+    kv_nope = tl.load(all_compressed_kv + off_kv_nope)
+    kv_rope = tl.load(all_compressed_kv + off_kv_rope)
+    off_nope = (offs_m + cur_batch_start_loc)[:, None] * stride_nope_s + offs_nope_d[None, :]
+    off_rope = (offs_m + cur_batch_start_loc)[:, None] * stride_rope_s + offs_rope_d[None, :]
+    nope_ptrs = sampled_compressed_kv_nope + off_nope
+    rope_ptrs = sampled_k_rope + off_rope
+    tl.store(nope_ptrs, kv_nope)
+    tl.store(rope_ptrs, kv_rope)
     return
 
 
 @torch.no_grad()
 def sample_kv(
-    kv_input,
-    kv_nope,
-    kv_rope,
-    b_req_idx,
-    max_value_in_b_seq_len,
-    b_seq_len,
-    req_to_token_indexs,
-    b_kv_start_loc,
-    kv_scale=None,
-    k_scale=None,
+    all_compressed_kv: torch.Tensor,
+    sampled_compressed_kv_nope: torch.Tensor,
+    sampled_k_rope: torch.Tensor,
+    b_req_idx: torch.Tensor,
+    req_to_token_indexs: torch.Tensor,
+    b_seq_len: torch.Tensor,
+    b_kv_start_loc: torch.Tensor,
+    max_kv_seq_len: int,
 ):
-    BLOCK = 128 if not is_tesla() else 64
-
-    nope_dim = kv_nope.shape[-1]
-    rope_dim = kv_rope.shape[-1]
-    if nope_dim >= 512:
-        BLOCK = 64 if not is_tesla() else 32
-    else:
-        BLOCK = 128 if not is_tesla() else 64
-
+    nope_dim = sampled_compressed_kv_nope.shape[-1]
+    rope_dim = sampled_k_rope.shape[-1]
+    assert rope_dim == 64
     batch = b_seq_len.shape[0]
 
-    max_input_len = max_value_in_b_seq_len
+    BLOCK = 64 if not is_tesla() else 32
+    num_warps = 8
     grid = (
         batch,
-        triton.cdiv(max_input_len, BLOCK),
+        triton.cdiv(max_kv_seq_len, BLOCK),
     )
-    num_warps = 4 if nope_dim <= 64 else 8
+
+    all_compressed_kv = all_compressed_kv.view(all_compressed_kv.shape[0], all_compressed_kv.shape[2])
+    sampled_compressed_kv_nope = sampled_compressed_kv_nope.view(sampled_compressed_kv_nope.shape[0], nope_dim)
+    sampled_k_rope = sampled_k_rope.view(sampled_k_rope.shape[0], rope_dim)
+    assert triton.next_power_of_2(nope_dim) == nope_dim
+    assert triton.next_power_of_2(rope_dim) == rope_dim
 
     _sample_kv_kernel[grid](
-        kv_input,
-        kv_scale,
-        kv_nope,
-        kv_rope,
-        k_scale,
-        b_kv_start_loc,
-        b_seq_len,
-        req_to_token_indexs,
-        b_req_idx,
-        kv_input.stride(0),
-        kv_scale.stride(0) if kv_scale is not None else 0,
-        kv_nope.stride(0),
-        kv_rope.stride(0),
-        req_to_token_indexs.stride(0),
-        HAS_SCALE=kv_scale is not None,
-        BLOCK_M=BLOCK,
-        BLOCK_DMODEL=nope_dim,
-        BLOCK_ROPE_DMODEL=rope_dim,
+        all_compressed_kv=all_compressed_kv,
+        stride_all_s=all_compressed_kv.stride(0),
+        stride_all_d=all_compressed_kv.stride(1),
+        sampled_compressed_kv_nope=sampled_compressed_kv_nope,
+        stride_nope_s=sampled_compressed_kv_nope.stride(0),
+        stride_nope_d=sampled_compressed_kv_nope.stride(1),
+        sampled_k_rope=sampled_k_rope,
+        stride_rope_s=sampled_k_rope.stride(0),
+        stride_rope_d=sampled_k_rope.stride(1),
+        b_kv_start_loc=b_kv_start_loc,
+        b_seq_len=b_seq_len,
+        req_to_token_indexs=req_to_token_indexs,
+        stride_req_to_tokens_b=req_to_token_indexs.stride(0),
+        b_req_idx=b_req_idx,
+        BLOCK_SEQ=BLOCK,
+        BLOCK_NOPE_DIM=nope_dim,
+        BLOCK_ROPE_DIM=rope_dim,
         num_warps=num_warps,
         num_stages=1,
     )
diff --git a/lightllm/models/deepseek_mtp/layer_infer/pre_layer_infer.py b/lightllm/models/deepseek_mtp/layer_infer/pre_layer_infer.py
index 26bfc865e4..7e12245587 100644
--- a/lightllm/models/deepseek_mtp/layer_infer/pre_layer_infer.py
+++ b/lightllm/models/deepseek_mtp/layer_infer/pre_layer_infer.py
@@ -8,8 +8,8 @@
 class Deepseek3MTPPreLayerInfer(LlamaPreLayerInfer):
     """ """
 
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         self.eps_ = network_config["rms_norm_eps"]
         self.hidden_size = network_config["hidden_size"]
         return
@@ -22,16 +22,8 @@ def _mtp_context_forward(
             input_embdings.shape[0] == tgt_embdings.shape[0]
         ), f"shape {input_embdings.shape} != shape {tgt_embdings.shape}"
 
-        layer_weight.enorm_weight_.rmsnorm_forward(
-            input=input_embdings,
-            eps=self.eps_,
-            out=input_embdings,
-        )
-        layer_weight.hnorm_weight_.rmsnorm_forward(
-            input=tgt_embdings,
-            eps=self.eps_,
-            out=tgt_embdings,
-        )
+        layer_weight.enorm_weight_(input=input_embdings, eps=self.eps_, out=input_embdings)
+        layer_weight.hnorm_weight_(input=tgt_embdings, eps=self.eps_, out=tgt_embdings)
         cat_embdings = torch.cat((input_embdings, tgt_embdings), dim=-1)
 
         ans_logics = layer_weight.eh_proj_weight_.mm(cat_embdings)
@@ -43,16 +35,8 @@ def _mtp_token_forward(
         tgt_embdings = infer_state.mtp_draft_input_hiddens
         assert input_embdings.shape[0] == tgt_embdings.shape[0]
 
-        layer_weight.enorm_weight_.rmsnorm_forward(
-            input=input_embdings,
-            eps=self.eps_,
-            out=input_embdings,
-        )
-        layer_weight.hnorm_weight_.rmsnorm_forward(
-            input=tgt_embdings,
-            eps=self.eps_,
-            out=tgt_embdings,
-        )
+        layer_weight.enorm_weight_(input=input_embdings, eps=self.eps_, out=input_embdings)
+        layer_weight.hnorm_weight_(input=tgt_embdings, eps=self.eps_, out=tgt_embdings)
         cat_embdings = torch.cat((input_embdings, tgt_embdings), dim=-1)
 
         ans_logics = layer_weight.eh_proj_weight_.mm(cat_embdings)
diff --git a/lightllm/models/deepseek_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/deepseek_mtp/layer_weights/pre_and_post_layer_weight.py
index 4a5bf2e961..719c80c27f 100644
--- a/lightllm/models/deepseek_mtp/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/deepseek_mtp/layer_weights/pre_and_post_layer_weight.py
@@ -2,15 +2,16 @@
 from lightllm.common.basemodel.layer_weights.meta_weights import (
     EmbeddingWeight,
     LMHeadWeight,
-    NoTpNormWeight,
+    RMSNormWeight,
     ROWMMWeight,
 )
 
 
 class Deepseek3MTPPreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
 
+        hidden_size = network_config["hidden_size"]
         self.eh_proj_weight_ = ROWMMWeight(
             weight_names="model.layers.0.eh_proj.weight",
             data_type=self.data_type_,
@@ -18,17 +19,20 @@ def __init__(self, data_type, network_config, mode):
             tp_rank=0,
             tp_world_size=1,
         )
-        self.enorm_weight_ = NoTpNormWeight(
+        self.enorm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="model.layers.0.enorm.weight",
             data_type=self.data_type_,
             bias_name=None,
         )
-        self.hnorm_weight_ = NoTpNormWeight(
+        self.hnorm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="model.layers.0.hnorm.weight",
             data_type=self.data_type_,
             bias_name=None,
         )
-        self.final_norm_weight_ = NoTpNormWeight(
+        self.final_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="model.layers.0.shared_head.norm.weight",
             data_type=self.data_type_,
             bias_name=None,
diff --git a/lightllm/models/gemma3/layer_infer/post_layer_infer.py b/lightllm/models/gemma3/layer_infer/post_layer_infer.py
index 22dc595059..57b21844ec 100644
--- a/lightllm/models/gemma3/layer_infer/post_layer_infer.py
+++ b/lightllm/models/gemma3/layer_infer/post_layer_infer.py
@@ -4,7 +4,7 @@
 class Gemma3PostLayerInfer(LlamaPostLayerInfer):
     """ """
 
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         self.eps_ = 1e-6
         return
diff --git a/lightllm/models/gemma3/layer_infer/pre_layer_infer.py b/lightllm/models/gemma3/layer_infer/pre_layer_infer.py
index dc8a46ad91..3543786f69 100644
--- a/lightllm/models/gemma3/layer_infer/pre_layer_infer.py
+++ b/lightllm/models/gemma3/layer_infer/pre_layer_infer.py
@@ -5,8 +5,8 @@
 
 
 class Gemma3PreLayerInfer(LlamaMultimodalPreLayerInfer):
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         self.embed_scale = torch.tensor(network_config["hidden_size"] ** 0.5, dtype=torch.float32)
         self.boi_token_index: int = 255_999
         self.eoi_token_index: int = 256_000
diff --git a/lightllm/models/gemma3/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma3/layer_infer/transformer_layer_infer.py
index d4bd8c3fa6..183c4d8d45 100644
--- a/lightllm/models/gemma3/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/gemma3/layer_infer/transformer_layer_infer.py
@@ -1,12 +1,6 @@
 import torch
-import torch.functional as F
 import torch.distributed as dist
 import torch.nn as nn
-import numpy as np
-from typing import Tuple
-from functools import partial
-import triton
-
 from lightllm.common.basemodel.infer_struct import InferStateInfo
 from lightllm.distributed import all_reduce
 from lightllm.models.gemma3.layer_weights.transformer_layer_weight import Gemma3TransformerLayerWeight
@@ -18,8 +12,8 @@
 class Gemma3TransformerLayerInfer(LlamaTransformerLayerInfer):
     """ """
 
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         self.tp_k_head_num_ = network_config["num_key_value_heads"]
         self.tp_v_head_num_ = network_config["num_key_value_heads"]
         self.eps_ = 1e-6
@@ -43,14 +37,10 @@ def _get_qkv(
         q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
         k = cache_kv[:, 0 : self.tp_k_head_num_, :]
 
-        q = layer_weight.q_norm_weight_.rmsnorm_forward(
-            input=q.float(), eps=self.eps_, alloc_func=self.alloc_tensor
-        ).to(cache_kv.dtype)
+        q = layer_weight.q_norm_weight_(input=q.float(), eps=self.eps_, alloc_func=self.alloc_tensor).to(cache_kv.dtype)
 
-        cache_kv[:, 0 : self.tp_k_head_num_, :] = layer_weight.k_norm_weight_.rmsnorm_forward(
-            input=k.float(),
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
+        cache_kv[:, 0 : self.tp_k_head_num_, :] = layer_weight.k_norm_weight_(
+            input=k.float(), eps=self.eps_, alloc_func=self.alloc_tensor
         ).to(cache_kv.dtype)
 
         is_sliding = bool((self.layer_num_ + 1) % self.sliding_window_pattern)
@@ -98,7 +88,7 @@ def context_forward(self, input_embdings, infer_state: InferStateInfo, layer_wei
         input_embdings.add_(o.view(-1, self.embed_dim_))
         o = None
 
-        input1 = layer_weight.pre_feedforward_layernorm_weight_.rmsnorm_forward(
+        input1 = layer_weight.pre_feedforward_layernorm_weight_(
             input=input_embdings.float(), eps=self.eps_, alloc_func=self.alloc_tensor
         ).to(torch.bfloat16)
 
@@ -107,10 +97,8 @@ def context_forward(self, input_embdings, infer_state: InferStateInfo, layer_wei
         if self.tp_world_size_ > 1:
             all_reduce(ffn_out, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False)
 
-        ffn_out = layer_weight.post_feedforward_layernorm_weight_.rmsnorm_forward(
-            input=ffn_out.float(),
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
+        ffn_out = layer_weight.post_feedforward_layernorm_weight_(
+            input=ffn_out.float(), eps=self.eps_, alloc_func=self.alloc_tensor
         ).to(torch.bfloat16)
 
         input_embdings.add_(ffn_out.view(-1, self.embed_dim_))
@@ -133,7 +121,7 @@ def token_forward(self, input_embdings, infer_state: InferStateInfo, layer_weigh
         input_embdings.add_(o.view(-1, self.embed_dim_))
         o = None
 
-        input1 = layer_weight.pre_feedforward_layernorm_weight_.rmsnorm_forward(
+        input1 = layer_weight.pre_feedforward_layernorm_weight_(
             input=input_embdings.float(), eps=self.eps_, alloc_func=self.alloc_tensor
         ).to(torch.bfloat16)
 
@@ -142,10 +130,8 @@ def token_forward(self, input_embdings, infer_state: InferStateInfo, layer_weigh
         if self.tp_world_size_ > 1:
             all_reduce(ffn_out, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False)
 
-        ffn_out = layer_weight.post_feedforward_layernorm_weight_.rmsnorm_forward(
-            input=ffn_out.float(),
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
+        ffn_out = layer_weight.post_feedforward_layernorm_weight_(
+            input=ffn_out.float(), eps=self.eps_, alloc_func=self.alloc_tensor
         ).to(torch.bfloat16)
 
         input_embdings.add_(ffn_out.view(-1, self.embed_dim_))
diff --git a/lightllm/models/gemma3/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/gemma3/layer_weights/pre_and_post_layer_weight.py
index 17e65268cc..7ae0fbcca3 100644
--- a/lightllm/models/gemma3/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/gemma3/layer_weights/pre_and_post_layer_weight.py
@@ -3,16 +3,21 @@
 
 
 class Gemma3PreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
+        hidden_size = network_config["hidden_size"]
+        vocab_size = network_config["vocab_size"]
 
         self.wte_weight_ = EmbeddingWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
             weight_name="language_model.model.embed_tokens.weight",
             data_type=self.data_type_,
         )
         self.lm_head_weight_ = self.wte_weight_
 
         self.final_norm_weight_ = NoTpGEMMANormWeight(
+            dim=hidden_size,
             weight_name="language_model.model.norm.weight",
             data_type=self.data_type_,
             bias_name=None,
diff --git a/lightllm/models/gemma3/layer_weights/transformer_layer_weight.py b/lightllm/models/gemma3/layer_weights/transformer_layer_weight.py
index 1e7ceeb42a..a4340a17a7 100644
--- a/lightllm/models/gemma3/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/gemma3/layer_weights/transformer_layer_weight.py
@@ -9,10 +9,9 @@ def __init__(
         layer_num,
         data_type,
         network_config,
-        mode=[],
         quant_cfg=None,
     ):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _init_weight_names(self):
@@ -26,51 +25,63 @@ def _init_weight_names(self):
 
     def _init_ffn(self):
         self.gate_proj = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_inter],
             weight_names=self._gate_weight_name,
             data_type=self.data_type_,
             bias_names=self._gate_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="gate_proj",
+            quant_method=self.get_quant_method("gate_proj"),
         )
         self.up_proj = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_inter],
             weight_names=self._up_weight_name,
             data_type=self.data_type_,
             bias_names=self._up_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="up_proj",
+            quant_method=self.get_quant_method("up_proj"),
         )
         super()._init_ffn()
 
     def _init_qkv(self):
+        kv_out_dim = self.k_head_num_ * self.head_dim
         self.k_proj = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[kv_out_dim],
             weight_names=self._k_weight_name,
             data_type=self.data_type_,
             bias_names=self._k_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="k_proj",
+            quant_method=self.get_quant_method("k_proj"),
         )
         self.v_proj = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[kv_out_dim],
             weight_names=self._v_weight_name,
             data_type=self.data_type_,
             bias_names=self._v_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="v_proj",
+            quant_method=self.get_quant_method("v_proj"),
         )
         super()._init_qkv()
 
     def _init_norm(self):
         super()._init_norm()
-        self.k_norm_weight_ = NoTpGEMMANormWeight(self._k_norm_weight_name, self.data_type_, bias_name=None)
-        self.q_norm_weight_ = NoTpGEMMANormWeight(self._q_norm_weight_name, self.data_type_, bias_name=None)
+
+        self.k_norm_weight_ = NoTpGEMMANormWeight(
+            dim=self.head_dim_, weight_name=self._k_norm_weight_name, data_type=self.data_type_, bias_name=None
+        )
+        self.q_norm_weight_ = NoTpGEMMANormWeight(
+            dim=self.head_dim_, weight_name=self._q_norm_weight_name, data_type=self.data_type_, bias_name=None
+        )
         self.pre_feedforward_layernorm_weight_ = NoTpGEMMANormWeight(
-            self._pre_feedforward_layernorm_name, self.data_type_, bias_name=None
+            dim=self.n_embed,
+            weight_name=self._pre_feedforward_layernorm_name,
+            data_type=self.data_type_,
+            bias_name=None,
         )
         self.post_feedforward_layernorm_weight_ = NoTpGEMMANormWeight(
-            self._post_feedforward_layernorm_name, self.data_type_, bias_name=None
+            dim=self.n_embed,
+            weight_name=self._post_feedforward_layernorm_name,
+            data_type=self.data_type_,
+            bias_name=None,
         )
 
     def load_hf_weights(self, weights):
diff --git a/lightllm/models/gemma3/model.py b/lightllm/models/gemma3/model.py
index dc4f03b7e1..9931c31713 100644
--- a/lightllm/models/gemma3/model.py
+++ b/lightllm/models/gemma3/model.py
@@ -1,7 +1,5 @@
 import os
-import re
 import json
-import numpy as np
 import torch
 from lightllm.models.registry import ModelRegistry
 from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer
@@ -14,8 +12,6 @@
 from lightllm.models.gemma3.layer_weights.pre_and_post_layer_weight import Gemma3PreAndPostLayerWeight
 from lightllm.models.gemma3.layer_weights.transformer_layer_weight import Gemma3TransformerLayerWeight
 from lightllm.models.llama.model import LlamaTpPartModel
-from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer
-from lightllm.models.llava.layer_weights.pre_and_post_layer_weight import LlavaPreAndPostLayerWeight
 from lightllm.server.multimodal_params import AudioItem, MultimodalParams, ImageItem
 from lightllm.server.core.objs import SamplingParams
 from lightllm.common.build_utils import repair_config
diff --git a/lightllm/models/gemma_2b/layer_infer/pre_layer_infer.py b/lightllm/models/gemma_2b/layer_infer/pre_layer_infer.py
index ce9737820e..e21788d762 100644
--- a/lightllm/models/gemma_2b/layer_infer/pre_layer_infer.py
+++ b/lightllm/models/gemma_2b/layer_infer/pre_layer_infer.py
@@ -11,8 +11,8 @@
 class Gemma_2bPreLayerInfer(PreLayerInferTpl):
     """ """
 
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         tp_vob_ids = np.linspace(0, network_config["vocab_size"], self.tp_world_size_ + 1, dtype=np.int64)
         self.vob_start_id_, self.vob_end_id_ = int(tp_vob_ids[self.tp_rank_]), int(tp_vob_ids[self.tp_rank_ + 1])
         self.normfactor = network_config["hidden_size"] ** 0.5
@@ -22,20 +22,14 @@ def _norm(self, input, infer_state, layer_weight: Gemma_2bPreAndPostLayerWeight)
         return input * self.normfactor
 
     def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: Gemma_2bPreAndPostLayerWeight):
-        input_embdings = layer_weight.wte_weight_.embedding(
-            input_ids=input_ids,
-            alloc_func=self.alloc_tensor,
-        )
+        input_embdings = layer_weight.wte_weight_(input_ids=input_ids, alloc_func=self.alloc_tensor)
         if self.tp_world_size_ > 1:
             all_reduce(input_embdings, group=infer_state.dist_group, op=dist.ReduceOp.SUM, async_op=False)
         input_embdings = self._norm(input_embdings, infer_state, layer_weight)
         return input_embdings
 
     def token_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: Gemma_2bPreAndPostLayerWeight):
-        input_embdings = layer_weight.wte_weight_.embedding(
-            input_ids=input_ids,
-            alloc_func=self.alloc_tensor,
-        )
+        input_embdings = layer_weight.wte_weight_(input_ids=input_ids, alloc_func=self.alloc_tensor)
         if self.tp_world_size_ > 1:
             all_reduce(input_embdings, group=infer_state.dist_group, op=dist.ReduceOp.SUM, async_op=False)
         input_embdings = self._norm(input_embdings, infer_state, layer_weight)
diff --git a/lightllm/models/gemma_2b/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma_2b/layer_infer/transformer_layer_infer.py
index 35ddaef343..2ed325659d 100644
--- a/lightllm/models/gemma_2b/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/gemma_2b/layer_infer/transformer_layer_infer.py
@@ -16,8 +16,8 @@
 class Gemma_2bTransformerLayerInfer(LlamaTransformerLayerInfer):
     """ """
 
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         self.tp_k_head_num_ = network_config["num_key_value_heads"]  # [SYM] always == 1
         self.tp_v_head_num_ = network_config["num_key_value_heads"]
         return
diff --git a/lightllm/models/gemma_2b/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/gemma_2b/layer_weights/pre_and_post_layer_weight.py
index d5d0438fa3..23ae50f096 100644
--- a/lightllm/models/gemma_2b/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/gemma_2b/layer_weights/pre_and_post_layer_weight.py
@@ -3,16 +3,21 @@
 
 
 class Gemma_2bPreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
+        hidden_size = network_config["hidden_size"]
+        vocab_size = network_config["vocab_size"]
 
         self.wte_weight_ = EmbeddingWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
             weight_name="model.embed_tokens.weight",
             data_type=self.data_type_,
         )
         self.lm_head_weight_ = self.wte_weight_
 
         self.final_norm_weight_ = NoTpGEMMANormWeight(
+            dim=hidden_size,
             weight_name="model.norm.weight",
             data_type=self.data_type_,
             bias_name=None,
diff --git a/lightllm/models/gemma_2b/layer_weights/transformer_layer_weight.py b/lightllm/models/gemma_2b/layer_weights/transformer_layer_weight.py
index 1916bd095c..19951b9900 100644
--- a/lightllm/models/gemma_2b/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/gemma_2b/layer_weights/transformer_layer_weight.py
@@ -6,28 +6,35 @@
 
 
 class Gemma_2bTransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _init_qkv(self):
+        in_dim = self.n_embed
+        q_out_dim = self.q_head_num_ * self.head_dim
+        kv_out_dim = self.k_head_num_ * self.head_dim
         self.q_proj = ROWMMWeight(
+            in_dim=in_dim,
+            out_dims=[q_out_dim],
             weight_names=self._q_weight_name,
             data_type=self.data_type_,
             bias_names=self._q_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="q_proj",
+            quant_method=self.get_quant_method("q_proj"),
         )
         self.kv_proj = ROWMMWeight(
+            in_dim=in_dim,
+            out_dims=[kv_out_dim, kv_out_dim],
             weight_names=[self._k_weight_name, self._v_weight_name],
             data_type=self.data_type_,
             bias_names=[self._k_bias_name, self._v_bias_name],
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="kv_proj",
+            quant_method=self.get_quant_method("kv_proj"),
         )
 
     def _init_norm(self):
-        self.att_norm_weight_ = NoTpGEMMANormWeight(self._att_norm_weight_name, self.data_type_)
-        self.ffn_norm_weight_ = NoTpGEMMANormWeight(self._ffn_norm_weight_name, self.data_type_)
+        self.att_norm_weight_ = NoTpGEMMANormWeight(
+            dim=self.n_embed, weight_name=self._att_norm_weight_name, data_type=self.data_type_, bias_name=None
+        )
+        self.ffn_norm_weight_ = NoTpGEMMANormWeight(
+            dim=self.n_embed, weight_name=self._ffn_norm_weight_name, data_type=self.data_type_, bias_name=None
+        )
diff --git a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py
index 93cd7413ba..d80eefd16e 100644
--- a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py
@@ -1,22 +1,15 @@
 import torch
-from torch import nn
-from torch.nn import functional as F
-import numpy as np
-from functools import partial
-from typing import Optional
-
 from lightllm.models.gpt_oss.layer_weights.transformer_layer_weight import GptOssTransformerLayerWeight
-from lightllm.models.llama.flashattention_infer_struct import FlashAttentionStateInfo
-from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
-from lightllm.utils.sgl_utils import flash_attn_with_kvcache
+from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer, LlamaInferStateInfo
+from lightllm.common.basemodel.attention.base_att import AttControl
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
 
 
 class GptOssTransformerLayerInfer(LlamaTransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         self.hidden_size = self.network_config_["hidden_size"]
         self.alpha = 1.702
         self.limit = 7.0
@@ -24,22 +17,17 @@ def __init__(self, layer_num, network_config, mode=[]):
         self.sliding_window = network_config["sliding_window"]
         self.head_dim_ = network_config["head_dim"]
 
-    def _bind_attention(self):
-        self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-        self._context_attention_kernel = self._context_sliding_attention_flashattention
-        self._token_attention_kernel = self._token_sliding_attention_flashattention
-
     def _bind_norm(self):
         self._att_norm = self._att_norm
         self._ffn_norm = self._ffn_norm
         return
 
-    def _att_norm(self, input, infer_state, layer_weight) -> torch.Tensor:
+    def _att_norm(self, input, infer_state, layer_weight: GptOssTransformerLayerWeight) -> torch.Tensor:
         out = self.alloc_tensor(input.shape, input.dtype)
         out = self._gpt_oss_rmsnorm(input, weight=layer_weight.att_norm_weight_.weight, eps=self.eps_)
         return out
 
-    def _ffn_norm(self, input, infer_state, layer_weight) -> torch.Tensor:
+    def _ffn_norm(self, input, infer_state, layer_weight: GptOssTransformerLayerWeight) -> torch.Tensor:
         out = self.alloc_tensor(input.shape, input.dtype)
         out = self._gpt_oss_rmsnorm(input, weight=layer_weight.ffn_norm_weight_.weight, eps=self.eps_)
         return out
@@ -51,9 +39,7 @@ def _gpt_oss_rmsnorm(self, hidden_states, weight, eps=1e-6):
         hidden_states = hidden_states * torch.rsqrt(variance + eps)
         return (weight * hidden_states).to(input_dtype)  # main diff with Llama
 
-    def _ffn(
-        self, input, infer_state: FlashAttentionStateInfo, layer_weight: GptOssTransformerLayerWeight
-    ) -> torch.Tensor:
+    def _ffn(self, input, infer_state, layer_weight: GptOssTransformerLayerWeight) -> torch.Tensor:
         hidden_states = input.view(-1, self.embed_dim_)
         num_tokens, hidden_dim = hidden_states.shape
         router_logits = layer_weight.moe_gate.mm(hidden_states)
@@ -68,78 +54,61 @@ def _ffn(
         )
         return hidden_states.view(num_tokens, hidden_dim)
 
-    def _context_sliding_attention_flashattention(
-        self, q, kv, infer_state: FlashAttentionStateInfo, layer_weight: GptOssTransformerLayerWeight, out=None
+    def _context_attention_kernel(
+        self,
+        q: torch.Tensor,
+        kv,
+        infer_state: LlamaInferStateInfo,
+        layer_weight: GptOssTransformerLayerWeight,
+        out=None,
     ):
         if self.network_config_["layer_types"][self.layer_num_] == "sliding_attention":
             window_size = (self.sliding_window - 1, self.sliding_window - 1)
+            use_sliding_window = True
         else:
             window_size = (-1, -1)
+            use_sliding_window = False
 
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :].reshape(
-            -1, 1, self.tp_k_head_num_, self.head_dim_
+        _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_)
+        _q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
+        o_tensor = infer_state.prefill_att_state.prefill_att(
+            q=_q,
+            k=_k,
+            v=_v,
+            att_control=AttControl(
+                use_sliding_window=use_sliding_window,
+                sliding_window=window_size,
+                use_att_sink=True,
+                sink_weight=layer_weight.attn_sinks.weight,
+            ),
+            alloc_func=self.alloc_tensor,
         )
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ].reshape(-1, 1, self.tp_v_head_num_, self.head_dim_)
-        q = q.reshape(-1, self.tp_q_head_num_, self.head_dim_)
-        k_descale, v_descale = None, None  # disable quantization
-        Lq = q.shape[-1]
-        sm_scale = 1.0 / (Lq ** 0.5)
-        o = flash_attn_with_kvcache(
-            q=q,
-            k_cache=cache_k,
-            v_cache=cache_v,
-            page_table=infer_state.page_table,
-            cache_seqlens=infer_state.b_seq_len,
-            cu_seqlens_q=infer_state.cu_seqlens_q,
-            cu_seqlens_k_new=infer_state.cu_seqlens_k,
-            max_seqlen_q=infer_state.q_max_seq_len,
-            softmax_scale=sm_scale,
-            causal=True,
-            window_size=window_size,
-            softcap=0.0,
-            k_descale=k_descale,
-            v_descale=v_descale,
-            return_softmax_lse=False,
-            sinks=layer_weight.attn_sinks.weight,
-        )
-        return o
+        o_tensor = o_tensor.view(q.shape)
+        return o_tensor
 
-    def _token_sliding_attention_flashattention(
-        self, q, infer_state: FlashAttentionStateInfo, layer_weight: GptOssTransformerLayerWeight, out=None
+    def _token_attention_kernel(
+        self, q: torch.Tensor, infer_state: LlamaInferStateInfo, layer_weight: GptOssTransformerLayerWeight, out=None
     ):
         if self.network_config_["layer_types"][self.layer_num_] == "sliding_attention":
             window_size = (self.sliding_window - 1, self.sliding_window - 1)
+            use_sliding_window = True
         else:
             window_size = (-1, -1)
+            use_sliding_window = False
 
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :].reshape(
-            -1, 1, self.tp_k_head_num_, self.head_dim_
-        )
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ].reshape(-1, 1, self.tp_v_head_num_, self.head_dim_)
-        q = q.reshape(-1, self.tp_q_head_num_, self.head_dim_)
-        k_descale, v_descale = None, None  # disable quantization
-        Lq = q.shape[-1]
-        sm_scale = 1.0 / (Lq ** 0.5)
-        o = flash_attn_with_kvcache(
-            q=q,
-            k_cache=cache_k,
-            v_cache=cache_v,
-            page_table=infer_state.page_table,
-            cache_seqlens=infer_state.b_seq_len,
-            cu_seqlens_q=infer_state.cu_seqlens_q,
-            cu_seqlens_k_new=infer_state.cu_seqlens_k,
-            max_seqlen_q=1,
-            softmax_scale=sm_scale,
-            causal=True,
-            window_size=window_size,
-            softcap=0.0,
-            k_descale=k_descale,
-            v_descale=v_descale,
-            return_softmax_lse=False,
-            sinks=layer_weight.attn_sinks.weight,
+        _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_)
+        _q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
+        o_tensor = infer_state.decode_att_state.decode_att(
+            q=_q,
+            k=_k,
+            v=_v,
+            att_control=AttControl(
+                use_sliding_window=use_sliding_window,
+                sliding_window=window_size,
+                use_att_sink=True,
+                sink_weight=layer_weight.attn_sinks.weight,
+            ),
+            alloc_func=self.alloc_tensor,
         )
-        return o
+        o_tensor = o_tensor.view(q.shape)
+        return o_tensor
diff --git a/lightllm/models/gpt_oss/layer_weights/transformer_layer_weight.py b/lightllm/models/gpt_oss/layer_weights/transformer_layer_weight.py
index f6a841b1aa..e6d58c3b2f 100644
--- a/lightllm/models/gpt_oss/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/gpt_oss/layer_weights/transformer_layer_weight.py
@@ -2,11 +2,14 @@
 import torch
 import numpy as np
 
-from lightllm.common.basemodel.layer_weights.meta_weights.gpt_oss_fused_moe_weight_tp import GPTOSSFusedMoeWeightTP
+from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.gpt_oss_fused_moe_weight_tp import (
+    GPTOSSFusedMoeWeightTP,
+)
 from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight import ROWMMWeight
 from lightllm.common.basemodel.layer_weights.meta_weights import TpAttSinkWeight
 from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.envs_utils import get_env_start_args
 
 logger = init_logger(__name__)
 
@@ -17,24 +20,24 @@ def __init__(
         layer_num,
         data_type,
         network_config,
-        mode=[],
         quant_cfg=None,
     ):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _init_moe(self):
-        moe_mode = os.getenv("MOE_MODE", "TP")
+        enable_ep_moe = get_env_start_args().enable_ep_moe
         moe_intermediate_size = self.network_config_["intermediate_size"]
         n_routed_experts = self.network_config_["num_local_experts"]
-        assert moe_mode in ["TP"], "For now, GPT-OSS type model only support MOE TP mode."
+        assert not enable_ep_moe, "For now, GPT-OSS type model only support MOE TP mode."
 
         self.moe_gate = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[n_routed_experts],
             weight_names=self._router_weight_name,
             data_type=self.data_type_,
-            layer_num=self.layer_num_,
             bias_names=self._router_bias_name,
-            name="moe_gate",
+            quant_method=self.get_quant_method("moe_gate"),
             tp_rank=0,
             tp_world_size=1,
         )
diff --git a/lightllm/models/gpt_oss/model.py b/lightllm/models/gpt_oss/model.py
index 34a017b316..dc5f2abdfe 100644
--- a/lightllm/models/gpt_oss/model.py
+++ b/lightllm/models/gpt_oss/model.py
@@ -19,4 +19,9 @@ class GptOssTpPartModel(LlamaTpPartModel):
 
     def __init__(self, kvargs):
         super().__init__(kvargs)
-        assert get_env_start_args().enable_fa3, "For now GPT-OSS type model only support flashattention-3"
+        assert (
+            get_env_start_args().llm_prefill_att_backend[0] == "fa3"
+        ), "For now GPT-OSS type model only support flashattention-3"
+        assert (
+            get_env_start_args().llm_decode_att_backend[0] == "fa3"
+        ), "For now GPT-OSS type model only support flashattention-3"
diff --git a/lightllm/models/internlm/layer_weights/transformer_layer_weight.py b/lightllm/models/internlm/layer_weights/transformer_layer_weight.py
index a2fc91dc46..6ef81122d3 100755
--- a/lightllm/models/internlm/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/internlm/layer_weights/transformer_layer_weight.py
@@ -1,13 +1,9 @@
-import torch
-import math
-import numpy as np
-
 from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
 
 
 class InternlmTransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _init_weight_names(self):
diff --git a/lightllm/models/internlm/model.py b/lightllm/models/internlm/model.py
index 78ac7117e1..50adbb3f9f 100644
--- a/lightllm/models/internlm/model.py
+++ b/lightllm/models/internlm/model.py
@@ -1,6 +1,3 @@
-import os
-import json
-import torch
 from lightllm.models.registry import ModelRegistry
 from lightllm.models.internlm.layer_weights.transformer_layer_weight import InternlmTransformerLayerWeight
 from lightllm.models.llama.model import LlamaTpPartModel
diff --git a/lightllm/models/internlm2/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/internlm2/layer_weights/pre_and_post_layer_weight.py
index b40330aa3d..3bb526c79b 100644
--- a/lightllm/models/internlm2/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/internlm2/layer_weights/pre_and_post_layer_weight.py
@@ -1,14 +1,26 @@
 from lightllm.common.basemodel import PreAndPostLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, NoTpNormWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, RMSNormWeight
 
 
 class Internlm2PreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
-        self.wte_weight_ = EmbeddingWeight(weight_name="model.tok_embeddings.weight", data_type=self.data_type_)
-        self.lm_head_weight_ = LMHeadWeight(weight_name="output.weight", data_type=self.data_type_)
-
-        self.final_norm_weight_ = NoTpNormWeight(
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
+        hidden_size = network_config["hidden_size"]
+        vocab_size = network_config["vocab_size"]
+        self.wte_weight_ = EmbeddingWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
+            weight_name="model.tok_embeddings.weight",
+            data_type=self.data_type_,
+        )
+        self.lm_head_weight_ = LMHeadWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
+            weight_name="output.weight",
+            data_type=self.data_type_,
+        )
+        self.final_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="model.norm.weight",
             data_type=self.data_type_,
         )
diff --git a/lightllm/models/internlm2/layer_weights/transformer_layer_weight.py b/lightllm/models/internlm2/layer_weights/transformer_layer_weight.py
index a675558632..a05e977f16 100755
--- a/lightllm/models/internlm2/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/internlm2/layer_weights/transformer_layer_weight.py
@@ -2,8 +2,8 @@
 
 
 class Internlm2TransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def load_hf_weights(self, weights):
diff --git a/lightllm/models/internlm2_reward/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/internlm2_reward/layer_weights/pre_and_post_layer_weight.py
index b20b9c4955..b526192125 100644
--- a/lightllm/models/internlm2_reward/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/internlm2_reward/layer_weights/pre_and_post_layer_weight.py
@@ -1,12 +1,15 @@
-import numpy as np
 from lightllm.common.basemodel import PreAndPostLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, NoTpNormWeight, ROWMMWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, RMSNormWeight, ROWMMWeight
 
 
 class Internlm2RewardPreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
+        hidden_size = network_config["hidden_size"]
+        vocab_size = network_config["vocab_size"]
         self.wte_weight_ = EmbeddingWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
             weight_name="model.tok_embeddings.weight",
             data_type=self.data_type_,
         )
@@ -17,7 +20,8 @@ def __init__(self, data_type, network_config, mode):
             tp_rank=0,
             tp_world_size=1,
         )
-        self.final_norm_weight_ = NoTpNormWeight(
+        self.final_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="model.norm.weight",
             data_type=self.data_type_,
         )
diff --git a/lightllm/models/internvl/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/internvl/layer_weights/pre_and_post_layer_weight.py
index 7d76d202ae..21a4c2e6b5 100644
--- a/lightllm/models/internvl/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/internvl/layer_weights/pre_and_post_layer_weight.py
@@ -13,8 +13,8 @@ def rename_weight_keys(weights):
 
 
 class InternVLPhi3PreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
         return
 
     def load_hf_weights(self, weights):
@@ -24,8 +24,8 @@ def load_hf_weights(self, weights):
 
 
 class InternVLInternlm2PreAndPostLayerWeight(Internlm2PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
         return
 
     def load_hf_weights(self, weights):
@@ -35,8 +35,8 @@ def load_hf_weights(self, weights):
 
 
 class InternVLLlamaPreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
         return
 
     def load_hf_weights(self, weights):
diff --git a/lightllm/models/internvl/model.py b/lightllm/models/internvl/model.py
index 6d264a4267..ccb76d3512 100644
--- a/lightllm/models/internvl/model.py
+++ b/lightllm/models/internvl/model.py
@@ -149,6 +149,10 @@ def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs):
                     raise ValueError("image token error")
             except ValueError:
                 break
+        if multimodal_params:
+            image_cnt = len(multimodal_params.images)
+            if image_cnt != image_id:
+                raise ValueError(image_cnt == image_id, f"invalid image tag num: {image_cnt} vs {image_id}!")
         input_ids.extend(origin_ids[start_idx:])
 
         # audio
@@ -174,6 +178,10 @@ def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs):
                     raise ValueError("audio token error")
             except ValueError:
                 break
+        if multimodal_params:
+            audio_cnt = len(multimodal_params.audios)
+            if audio_cnt != audio_id:
+                raise ValueError(audio_cnt == audio_id, f"invalid audio tag num: {audio_cnt} vs {audio_id}!")
         input_ids.extend(origin_ids[start_idx:])
         return input_ids
 
diff --git a/lightllm/models/llama/flashattention_infer_struct.py b/lightllm/models/llama/flashattention_infer_struct.py
deleted file mode 100644
index 9f71cbbc56..0000000000
--- a/lightllm/models/llama/flashattention_infer_struct.py
+++ /dev/null
@@ -1,106 +0,0 @@
-import os
-import torch
-import numpy as np
-import torch.distributed as dist
-from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.utils.envs_utils import get_env_start_args
-from lightllm.utils.dist_utils import get_current_device_id
-from lightllm.models.deepseek2.triton_kernel.repack_kv_index import repack_kv_index
-from lightllm.common.basemodel.batch_objs import ModelInput
-from lightllm.common.basemodel.triton_kernel.fa3_utils import page_table_copy
-
-
-class FlashAttentionStateInfo(LlamaInferStateInfo):
-    _shared_page_table_buffer = None
-
-    def __init__(self):
-        super().__init__()
-
-    @classmethod
-    def get_page_table_buffer(cls, graph_max_batch_size: int, max_seq_len: int):
-        if cls._shared_page_table_buffer is None:
-            cls._shared_page_table_buffer = [
-                torch.empty(graph_max_batch_size * max_seq_len, dtype=torch.int32).to(get_current_device_id()),
-                torch.empty(graph_max_batch_size * max_seq_len, dtype=torch.int32).to(get_current_device_id()),
-            ]
-        return cls._shared_page_table_buffer
-
-    def _init_flash_attention_state(self, model):
-        if self.is_prefill:
-            self.cu_seqlens_q = self.b1_cu_q_seq_len.int()
-            self.cu_seqlens_k = self.b1_cu_kv_seq_len.int()
-            self.page_table = torch.empty(
-                (self.batch_size, self.max_seq_len), dtype=torch.int32, device=self.input_ids.device
-            )
-            self.page_table.copy_(model.req_manager.req_to_token_indexs[self.b_req_idx, : self.max_seq_len])
-        else:
-            # Meta information of flashattention for decoding
-            self.cu_seqlens_q = self.b1_cu_q_seq_len.int()
-            self.cu_seqlens_k = self.b1_cu_kv_seq_len.int()
-            max_seq_len_k = self.max_kv_seq_len
-            args_mtp_step = get_env_start_args().mtp_step
-            att_batch_size = self.batch_size // (args_mtp_step + 1)
-            if self.batch_size <= model.graph_max_batch_size and self.max_len_in_batch <= model.graph_max_len_in_batch:
-                page_buffer = FlashAttentionStateInfo.get_page_table_buffer(
-                    model.graph_max_batch_size, model.graph_max_len_in_batch
-                )
-                self.page_table = page_buffer[self.microbatch_index][
-                    : att_batch_size * model.graph_max_len_in_batch
-                ].reshape(att_batch_size, model.graph_max_len_in_batch)
-            else:
-                self.page_table = torch.empty(
-                    (att_batch_size, self.max_len_in_batch), dtype=torch.int32, device=self.input_ids.device
-                )
-
-            page_table_copy(
-                page_table=self.page_table[:, :max_seq_len_k],
-                req_to_token_indexs=model.req_manager.req_to_token_indexs,
-                b_req_idx=self.b_req_idx[args_mtp_step :: (args_mtp_step + 1)],
-            )
-            if args_mtp_step > 0:
-                self.b_att_seq_len = self.b_seq_len[args_mtp_step :: (args_mtp_step + 1)].contiguous()
-            else:
-                self.b_att_seq_len = self.b_seq_len
-
-        if "offline_calibration_fp8kv" in model.mode:
-            if self.is_prefill:
-                device = self.input_ids.device
-                # q_scale和token_batch_ids在对q做per head量化使用，为了节省资源在推理外部初始化
-                self.q_scale = torch.empty(
-                    (self.batch_size, self.mem_manager.head_num), dtype=torch.float32, device=device
-                )
-                self.token_batch_ids = torch.repeat_interleave(
-                    torch.arange(self.batch_size, device=device), self.b_q_seq_len
-                )
-
-            offline_scales = self.mem_manager.scales
-            head_num = self.mem_manager.head_num
-            # 为了减少推理计算量，在推理外部初始化k_descale和v_descale
-            self.k_descale = (
-                offline_scales[:, :head_num]
-                .view(-1, 1, head_num)
-                .expand(offline_scales.shape[0], self.batch_size, head_num)
-                if offline_scales is not None
-                else torch.ones(
-                    (self.mem_manager.layer_num, self.batch_size, head_num),
-                    dtype=torch.float32,
-                    device=self.input_ids.device,
-                )
-            )
-            self.v_descale = (
-                offline_scales[:, head_num:]
-                .view(-1, 1, head_num)
-                .expand(offline_scales.shape[0], self.batch_size, head_num)
-                if offline_scales is not None
-                else torch.ones(
-                    (self.mem_manager.layer_num, self.batch_size, head_num),
-                    dtype=torch.float32,
-                    device=self.input_ids.device,
-                )
-            )
-        return
-
-    def init_some_extra_state(self, model):
-        super().init_some_extra_state(model)
-        self._init_flash_attention_state(model)
-        return
diff --git a/lightllm/models/llama/flashinfer_struct.py b/lightllm/models/llama/flashinfer_struct.py
deleted file mode 100644
index 7f9beac1db..0000000000
--- a/lightllm/models/llama/flashinfer_struct.py
+++ /dev/null
@@ -1,127 +0,0 @@
-import os
-import torch
-import numpy as np
-import torch.distributed as dist
-from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.utils.envs_utils import get_env_start_args
-from lightllm.models.deepseek2.triton_kernel.repack_kv_index import repack_kv_index
-
-
-class LlamaFlashInferStateInfo(LlamaInferStateInfo):
-    def __init__(self):
-        super().__init__()
-        self.prefill_wrapper = None
-        self.decode_wrapper = None
-        self.flashinfer_extra_state = None
-
-    def init_some_extra_state(self, model):
-        super().init_some_extra_state(model)
-        self.flashinfer_extra_state = model.flashinfer_extra_state
-
-        import flashinfer
-
-        if not self.is_prefill:
-            if get_env_start_args().enable_flashinfer_decode:
-                self.kv_last_page_len_buffer = torch.full(
-                    (self.batch_size,), 1, dtype=torch.int32, device=self.input_ids.device
-                )
-                if self.batch_size <= model.graph_max_batch_size:
-                    self.kv_indices = self.flashinfer_extra_state.kv_indices_buffer[self.microbatch_index][
-                        : self.batch_size * self.flashinfer_extra_state.max_seq_length
-                    ]
-                else:
-                    self.kv_indices = torch.empty(
-                        self.batch_size * self.flashinfer_extra_state.max_seq_length,
-                        dtype=torch.int32,
-                        device=self.input_ids.device,
-                    )
-
-                repack_kv_index(
-                    self.req_manager.req_to_token_indexs,
-                    self.b_req_idx,
-                    self.b_seq_len,
-                    self.b_start_loc,
-                    self.max_len_in_batch,
-                    self.kv_indices,
-                )
-                self.kv_starts = self.b1_cu_kv_seq_len.int()
-                if self.decode_wrapper is None:
-                    self.decode_wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(
-                        self.flashinfer_extra_state.workspace_buffer,
-                        "NHD",
-                        use_cuda_graph=True,
-                        use_tensor_cores=True,
-                        paged_kv_indptr_buffer=self.kv_starts,
-                        paged_kv_indices_buffer=self.kv_indices,
-                        paged_kv_last_page_len_buffer=self.kv_last_page_len_buffer,
-                    )
-                    self.decode_wrapper.plan(
-                        self.kv_starts,
-                        self.kv_indices,
-                        self.kv_last_page_len_buffer,
-                        self.flashinfer_extra_state.tp_q_head_num,
-                        self.flashinfer_extra_state.tp_kv_head_num,
-                        self.flashinfer_extra_state.head_dim,
-                        1,
-                        q_data_type=self.flashinfer_extra_state.q_data_type,
-                        kv_data_type=self.flashinfer_extra_state.kv_data_type,
-                        non_blocking=True,
-                    )
-        else:
-            if get_env_start_args().enable_flashinfer_prefill:
-                q_starts = self.b1_cu_q_seq_len.int()
-                kv_starts = self.b1_cu_kv_seq_len.int()
-                kv_last_page_len = torch.full((self.batch_size,), 1, dtype=torch.int32, device=self.input_ids.device)
-                kv_indices = torch.empty(
-                    self.batch_size * self.flashinfer_extra_state.max_seq_length,
-                    dtype=torch.int32,
-                    device=self.input_ids.device,
-                )
-                repack_kv_index(
-                    self.req_manager.req_to_token_indexs,
-                    self.b_req_idx,
-                    self.b_seq_len,
-                    kv_starts[:-1],
-                    self.max_kv_seq_len,
-                    kv_indices,
-                )
-                self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper(
-                    self.flashinfer_extra_state.workspace_buffer,
-                    qo_indptr_buf=q_starts,
-                    paged_kv_indptr_buf=kv_starts,
-                    paged_kv_indices_buf=kv_indices,
-                    paged_kv_last_page_len_buf=kv_last_page_len,
-                )
-                self.prefill_wrapper.plan(
-                    q_starts,
-                    kv_starts,
-                    kv_indices,
-                    kv_last_page_len,
-                    self.flashinfer_extra_state.tp_q_head_num,
-                    self.flashinfer_extra_state.tp_kv_head_num,
-                    self.flashinfer_extra_state.head_dim,
-                    1,
-                    causal=True,
-                    pos_encoding_mode="NONE",
-                    logits_soft_cap=0.0,
-                    q_data_type=self.flashinfer_extra_state.q_data_type,
-                    kv_data_type=self.flashinfer_extra_state.kv_data_type,
-                )
-        return
-
-    def copy_for_cuda_graph(self, new_infer_state):
-        super().copy_for_cuda_graph(new_infer_state)
-        if get_env_start_args().enable_flashinfer_decode and not self.is_prefill:
-            self.decode_wrapper.plan(
-                new_infer_state.kv_starts,
-                new_infer_state.kv_indices,
-                new_infer_state.kv_last_page_len_buffer,
-                new_infer_state.flashinfer_extra_state.tp_q_head_num,
-                new_infer_state.flashinfer_extra_state.tp_kv_head_num,
-                new_infer_state.flashinfer_extra_state.head_dim,
-                1,
-                q_data_type=new_infer_state.flashinfer_extra_state.q_data_type,
-                kv_data_type=new_infer_state.flashinfer_extra_state.kv_data_type,
-                non_blocking=True,
-            )
-        return
diff --git a/lightllm/models/llama/infer_struct.py b/lightllm/models/llama/infer_struct.py
index 3bba439767..fe6ca392a2 100644
--- a/lightllm/models/llama/infer_struct.py
+++ b/lightllm/models/llama/infer_struct.py
@@ -14,7 +14,6 @@ def init_some_extra_state(self, model):
         super().init_some_extra_state(model)
         if self.is_prefill:
             self.max_seq_len = self.max_kv_seq_len
-            self.q_max_seq_len = self.max_q_seq_len
             position_ids = self.position_ids
             self.position_cos = torch.index_select(model._cos_cached, 0, position_ids).view(position_ids.shape[0], -1)
             self.position_sin = torch.index_select(model._sin_cached, 0, position_ids).view(position_ids.shape[0], -1)
diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py
index 7c7b0ea39b..7714164151 100644
--- a/lightllm/models/llama/layer_infer/post_layer_infer.py
+++ b/lightllm/models/llama/layer_infer/post_layer_infer.py
@@ -13,13 +13,13 @@
 class LlamaPostLayerInfer(PostLayerInferTpl):
     """ """
 
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         self.eps_ = network_config["rms_norm_eps"]
         return
 
     def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> torch.Tensor:
-        return layer_weight.final_norm_weight_.rmsnorm_forward(input=input, eps=self.eps_, alloc_func=self.alloc_tensor)
+        return layer_weight.final_norm_weight_(input=input, eps=self.eps_, alloc_func=self.alloc_tensor)
 
     def _slice_get_last_input(self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo):
         embed_dim_ = input_embdings.shape[1]
@@ -66,7 +66,7 @@ def token_forward(
         input_embdings = None
         last_input = self._norm(last_input, infer_state, layer_weight)
         last_input = last_input.permute(1, 0).view(-1, token_num)
-        logic_batch = layer_weight.lm_head_weight_.lm_head(input=last_input, alloc_func=self.alloc_tensor)
+        logic_batch = layer_weight.lm_head_weight_(input=last_input, alloc_func=self.alloc_tensor)
         last_input = None
         vocab_size = layer_weight.lm_head_weight_.vocab_size
         if self.tp_world_size_ == 1:
diff --git a/lightllm/models/llama/layer_infer/pre_layer_infer.py b/lightllm/models/llama/layer_infer/pre_layer_infer.py
index ddb99e2627..63a2fe4d14 100644
--- a/lightllm/models/llama/layer_infer/pre_layer_infer.py
+++ b/lightllm/models/llama/layer_infer/pre_layer_infer.py
@@ -10,18 +10,18 @@
 class LlamaPreLayerInfer(PreLayerInferTpl):
     """ """
 
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         return
 
     def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight):
-        input_embdings = layer_weight.wte_weight_.embedding(input_ids=input_ids, alloc_func=self.alloc_tensor)
+        input_embdings = layer_weight.wte_weight_(input_ids=input_ids, alloc_func=self.alloc_tensor)
         if self.tp_world_size_ > 1:
             all_reduce(input_embdings, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False)
         return input_embdings
 
     def token_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight):
-        input_embdings = layer_weight.wte_weight_.embedding(input_ids=input_ids, alloc_func=self.alloc_tensor)
+        input_embdings = layer_weight.wte_weight_(input_ids=input_ids, alloc_func=self.alloc_tensor)
         if self.tp_world_size_ > 1:
             all_reduce(input_embdings, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False)
         return input_embdings
diff --git a/lightllm/models/llama/layer_infer/transformer_layer_infer.py b/lightllm/models/llama/layer_infer/transformer_layer_infer.py
old mode 100755
new mode 100644
index b08b2aa1fd..dc6f10be59
--- a/lightllm/models/llama/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/llama/layer_infer/transformer_layer_infer.py
@@ -1,51 +1,23 @@
 import torch
 import triton
-import torch.functional as F
 import torch.distributed as dist
-import numpy as np
-from typing import Tuple
 from functools import partial
-
 from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
-from lightllm.models.llama.triton_kernel.context_flashattention_nopad import (
-    context_attention_fwd,
-    context_attention_fwd_ppl_int8kv,
-)
-from lightllm.models.llama.triton_kernel.token_attention_nopad_att1 import token_att_fwd, token_att_fwd_int8k
-from lightllm.models.llama.triton_kernel.token_attention_nopad_softmax import token_softmax_fwd
-from lightllm.models.llama.triton_kernel.token_attention_nopad_reduceV import token_att_fwd2, token_att_fwd2_int8v
 from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd
-from lightllm.common.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
-
+from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.models.llama.flashattention_infer_struct import FlashAttentionStateInfo
-from lightllm.models.llama.flashinfer_struct import LlamaFlashInferStateInfo
-from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv, destindex_copy_quantize_kv
-from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8
 from lightllm.common.basemodel import TransformerLayerInferTpl
-from lightllm.models.llama.triton_kernel.ppl_quant_copy_kv import destindex_copy_dequantize_kv
 from lightllm.distributed.communication_op import all_gather_into_tensor, reduce_scatter_tensor
 from lightllm.utils.log_utils import init_logger
-from lightllm.utils.envs_utils import get_env_start_args
-from lightllm.utils.light_utils import HAS_LIGHTLLM_KERNEL, light_ops
-from lightllm.common.basemodel.triton_kernel.q_per_head_fp8_quant import q_per_head_fp8_quant
-from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops
-
-if HAS_VLLM:
-    scaled_fp8_quant = vllm_ops.scaled_fp8_quant
-else:
-    scaled_fp8_quant = None
 
 logger = init_logger(__name__)
 
-from lightllm.utils.sgl_utils import flash_attn_with_kvcache
-
 
 class LlamaTransformerLayerInfer(TransformerLayerInferTpl):
     """ """
 
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         self.eps_ = network_config["rms_norm_eps"]
         self.tp_q_head_num_ = network_config["num_attention_heads"] // self.tp_world_size_
         self.tp_k_head_num_ = max(network_config["num_key_value_heads"] // self.tp_world_size_, 1)
@@ -58,7 +30,6 @@ def __init__(self, layer_num, network_config, mode=[]):
 
     def _bind_func(self):
         self._bind_norm()
-        self._bind_attention()
         return
 
     def _bind_norm(self):
@@ -66,139 +37,44 @@ def _bind_norm(self):
         self._ffn_norm = partial(LlamaTransformerLayerInfer._ffn_norm, self)
         return
 
-    def _bind_attention(self):
-        if get_env_start_args().enable_fa3:
-            if "offline_calibration_fp8kv" in self.mode:
-                self._context_attention_kernel = partial(
-                    LlamaTransformerLayerInfer._context_attention_flashattention_fp8, self
-                )
-                self._token_attention_kernel = partial(
-                    LlamaTransformerLayerInfer._token_decode_attention_flashattention_fp8, self
-                )
-                self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_fp8kv, self)
-            elif "export_fp8kv_calibration" in self.mode:
-                self._context_attention_kernel = partial(
-                    LlamaTransformerLayerInfer._context_attention_flashattention, self
-                )
-                self._token_attention_kernel = partial(
-                    LlamaTransformerLayerInfer._token_decode_attention_flashattention, self
-                )
-                self._copy_kv_to_mem_cache = partial(
-                    LlamaTransformerLayerInfer._copy_kv_to_mem_cache_with_calibration, self
-                )
-            elif not self.mode:
-                self._context_attention_kernel = partial(
-                    LlamaTransformerLayerInfer._context_attention_flashattention, self
-                )
-                self._token_attention_kernel = partial(
-                    LlamaTransformerLayerInfer._token_decode_attention_flashattention, self
-                )
-                self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-            else:
-                raise Exception(f"Unsupported mode for fa3 backend: {self.mode}")
-            return
-        elif get_env_start_args().enable_flashinfer_prefill:
-            self._context_attention_kernel = partial(
-                LlamaTransformerLayerInfer._context_attention_flashinfer_kernel, self
-            )
-        else:
-            self._context_attention_kernel = partial(LlamaTransformerLayerInfer._context_attention_kernel, self)
-        if "ppl_int8kv" in self.mode:
-            self._token_attention_kernel = partial(LlamaTransformerLayerInfer._token_decode_attention_ppl_int8kv, self)
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_ppl_int8kv, self)
-            self._context_attention_kernel = partial(
-                LlamaTransformerLayerInfer._context_attention_kernel_ppl_int8kv, self
-            )
-        elif "ppl_int8kv_flashdecoding_diverse" in self.mode:
-            self._token_attention_kernel = partial(
-                LlamaTransformerLayerInfer._token_decode_attention_ppl_int8kv_flashdecoding_diverse, self
-            )
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_ppl_int8kv, self)
-            self._context_attention_kernel = partial(
-                LlamaTransformerLayerInfer._context_attention_kernel_ppl_int8kv, self
-            )
-        elif "ppl_int8kv_flashdecoding" in self.mode:
-            self._token_attention_kernel = partial(
-                LlamaTransformerLayerInfer._token_decode_attention_ppl_int8kv_flashdecoding, self
-            )
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_ppl_int8kv, self)
-            self._context_attention_kernel = partial(
-                LlamaTransformerLayerInfer._context_attention_kernel_ppl_int8kv, self
-            )
-        elif "ppl_int4kv_flashdecoding" in self.mode:
-            self._token_attention_kernel = partial(
-                LlamaTransformerLayerInfer._token_decode_attention_ppl_int4kv_flashdecoding, self
-            )
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_ppl_int4kv, self)
-        elif "ppl_fp16" in self.mode:
-            self._token_attention_kernel = partial(LlamaTransformerLayerInfer._token_decode_attention_ppl_fp16, self)
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-        elif "ppl_fp16_flashdecoding" in self.mode:
-            self._token_attention_kernel = partial(
-                LlamaTransformerLayerInfer._token_decode_attention_ppl_fp16_flashdecoding, self
-            )
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-        elif "triton_int8kv" in self.mode:
-            self._token_attention_kernel = partial(LlamaTransformerLayerInfer._token_decode_attention_int8kv, self)
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_int8kv, self)
-        elif "offline_calibration_fp8kv" in self.mode:
-            assert get_env_start_args().enable_flashinfer_prefill and get_env_start_args().enable_flashinfer_decode
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_fp8kv, self)
-            self._context_attention_kernel = partial(
-                LlamaTransformerLayerInfer._context_attention_flashinfer_kernel_fp8, self
-            )
-            self._token_attention_kernel = partial(
-                LlamaTransformerLayerInfer._token_decode_attention_flashinfer_fp8, self
-            )
-        elif "triton_flashdecoding" in self.mode:
-            self._token_attention_kernel = partial(
-                LlamaTransformerLayerInfer._token_decode_attention_flashdecoding, self
-            )
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-        elif "triton_gqa_attention" in self.mode:
-            self._token_attention_kernel = partial(LlamaTransformerLayerInfer._token_decode_gqa_attention_normal, self)
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-        elif "triton_gqa_flashdecoding" in self.mode:
-            self._token_attention_kernel = partial(
-                LlamaTransformerLayerInfer._token_decode_attention_gqa_flashdecoding, self
-            )
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-        elif "triton_gqa_flashdecoding_vsm" in self.mode:
-            self._token_attention_kernel = partial(
-                LlamaTransformerLayerInfer._token_decode_attention_gqa_flashdecoding_vsm, self
-            )
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-        elif "export_fp8kv_calibration" in self.mode:
-            self._token_attention_kernel = partial(LlamaTransformerLayerInfer._token_decode_attention_flashinfer, self)
-            self._copy_kv_to_mem_cache = partial(
-                LlamaTransformerLayerInfer._copy_kv_to_mem_cache_with_calibration, self
-            )
-        elif not self.mode:
-            if get_env_start_args().enable_flashinfer_decode:
-                self._token_attention_kernel = partial(
-                    LlamaTransformerLayerInfer._token_decode_attention_flashinfer, self
-                )
-            else:
-                self._token_attention_kernel = partial(LlamaTransformerLayerInfer._token_decode_attention_normal, self)
-            self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-        else:
-            raise Exception(f"Unsupported mode: {self.mode}")
+    def _context_attention_kernel(
+        self,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        infer_state: LlamaInferStateInfo,
+        layer_weight: LlamaTransformerLayerWeight,
+    ) -> torch.Tensor:
+        _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_)
+        _q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
+        o_tensor = infer_state.prefill_att_state.prefill_att(
+            q=_q,
+            k=_k,
+            v=_v,
+            alloc_func=self.alloc_tensor,
+        )
+        o_tensor = o_tensor.view(q.shape)
+        return o_tensor
 
-        return
+    def _token_attention_kernel(
+        self,
+        q: torch.Tensor,
+        infer_state: LlamaInferStateInfo,
+        layer_weight: LlamaTransformerLayerWeight,
+    ) -> torch.Tensor:
+        _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_)
+        _q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
+        o_tensor = infer_state.decode_att_state.decode_att(q=_q, k=_k, v=_v, alloc_func=self.alloc_tensor)
+        return o_tensor.view(q.shape)
 
     def _att_norm(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight
     ) -> torch.Tensor:
-        return layer_weight.att_norm_weight_.rmsnorm_forward(input=input, eps=self.eps_, alloc_func=self.alloc_tensor)
+        return layer_weight.att_norm_weight_(input=input, eps=self.eps_, alloc_func=self.alloc_tensor)
 
     def _ffn_norm(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight
     ) -> torch.Tensor:
-        return layer_weight.ffn_norm_weight_.rmsnorm_forward(
-            input=input,
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
-        )
+        return layer_weight.ffn_norm_weight_(input=input, eps=self.eps_, alloc_func=self.alloc_tensor)
 
     def _get_qkv(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight
@@ -241,163 +117,6 @@ def _tpsp_get_qkv(
 
         return q, cache_kv
 
-    def _context_attention_flashinfer_kernel_fp8(
-        self, q, kv, infer_state: LlamaFlashInferStateInfo, layer_weight, out=None
-    ) -> torch.Tensor:
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        kv = kv.unsqueeze(1)
-        k = kv[:, :, : self.tp_k_head_num_, :].view(torch.float8_e4m3fn)
-        v = kv[:, :, self.tp_k_head_num_ :, :].view(torch.float8_e4m3fn)
-        offline_scales = infer_state.mem_manager.scales_list
-        k_descale = offline_scales[self.layer_num_][0] if offline_scales is not None else None
-        v_descale = offline_scales[self.layer_num_][1] if offline_scales is not None else None
-        infer_state.prefill_wrapper.run(
-            q.view(q.shape[0], -1, self.head_dim_),
-            (k, v),
-            k_scale=k_descale,
-            v_scale=v_descale,
-            out=o_tensor.view(q.shape[0], -1, self.head_dim_),
-        )
-        return o_tensor
-
-    def _context_attention_flashinfer_kernel(
-        self, q, kv, infer_state: LlamaFlashInferStateInfo, layer_weight, out=None
-    ) -> torch.Tensor:
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        kv = kv.unsqueeze(1)
-        infer_state.prefill_wrapper.run(
-            q.view(q.shape[0], -1, self.head_dim_),
-            (kv[:, :, : self.tp_k_head_num_, :], kv[:, :, self.tp_k_head_num_ :, :]),
-            out=o_tensor.view(q.shape[0], -1, self.head_dim_),
-        )
-        return o_tensor
-
-    def _context_attention_kernel(
-        self, q, kv, infer_state: LlamaInferStateInfo, layer_weight, out=None
-    ) -> torch.Tensor:
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        context_attention_fwd(
-            q.view(-1, self.tp_q_head_num_, self.head_dim_),
-            kv[:, 0 : self.tp_k_head_num_, :],
-            kv[:, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :],
-            o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_),
-            infer_state.b_req_idx,
-            infer_state.b_start_loc,
-            infer_state.b_seq_len,
-            infer_state.b_ready_cache_len,
-            infer_state.max_len_in_batch,
-            infer_state.req_manager.req_to_token_indexs,
-        )
-        return o_tensor
-
-    def _context_attention_kernel_ppl_int8kv(
-        self, q, kv, infer_state: LlamaInferStateInfo, layer_weight, out=None
-    ) -> torch.Tensor:
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        batch_size = infer_state.b_seq_len.shape[0]
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        kv_scale = infer_state.mem_manager.scale_buffer[self.layer_num_]
-        max_seq_len = infer_state.max_seq_len
-        kv_dequant = self.alloc_tensor(
-            (batch_size, kv.shape[1], max_seq_len, kv.shape[2]), device=q.device, dtype=q.dtype
-        )
-        destindex_copy_dequantize_kv(
-            kv,
-            kv_scale,
-            infer_state.req_manager.req_to_token_indexs,
-            infer_state.b_seq_len,
-            infer_state.b_req_idx,
-            max_seq_len,
-            kv_dequant,
-        )
-        context_attention_fwd_ppl_int8kv(
-            q.view(-1, self.tp_q_head_num_, self.head_dim_),
-            kv_dequant[:, 0 : self.tp_k_head_num_, :, :],
-            kv_dequant[:, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :, :],
-            o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_),
-            infer_state.b_start_loc,
-            infer_state.b_seq_len,
-            infer_state.max_len_in_batch,
-            infer_state.b_ready_cache_len,
-        )
-        return o_tensor
-
-    def _context_attention_flashattention(self, q, kv, infer_state: FlashAttentionStateInfo, layer_weight, out=None):
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :].reshape(
-            -1, 1, self.tp_k_head_num_, self.head_dim_
-        )
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ].reshape(-1, 1, self.tp_v_head_num_, self.head_dim_)
-        q = q.reshape(-1, self.tp_q_head_num_, self.head_dim_)
-        k_descale, v_descale = None, None  # disable quantization
-        Lq = q.shape[-1]
-        sm_scale = 1.0 / (Lq ** 0.5)
-        o = flash_attn_with_kvcache(
-            q=q,
-            k_cache=cache_k,
-            v_cache=cache_v,
-            page_table=infer_state.page_table,
-            cache_seqlens=infer_state.b_seq_len,
-            cu_seqlens_q=infer_state.cu_seqlens_q,
-            cu_seqlens_k_new=infer_state.cu_seqlens_k,
-            max_seqlen_q=infer_state.q_max_seq_len,
-            softmax_scale=sm_scale,
-            causal=True,
-            window_size=(-1, -1),
-            softcap=0.0,
-            k_descale=k_descale,
-            v_descale=v_descale,
-            return_softmax_lse=False,
-        )
-        return o
-
-    def _context_attention_flashattention_fp8(
-        self, q, kv, infer_state: FlashAttentionStateInfo, layer_weight, out=None
-    ):
-        q, q_scale = q_per_head_fp8_quant(
-            q.view(q.shape[0], self.tp_k_head_num_, -1),
-            infer_state.b_seq_len,
-            infer_state.cu_seqlens_q,
-            infer_state.q_scale,
-            infer_state.token_batch_ids,
-        )
-        cache_k = (
-            (infer_state.mem_manager.kv_buffer[self.layer_num_][:, : self.tp_k_head_num_, :])
-            .reshape(-1, 1, self.tp_k_head_num_, self.head_dim_)
-            .view(torch.float8_e4m3fn)
-        )
-        cache_v = (
-            (
-                infer_state.mem_manager.kv_buffer[self.layer_num_][
-                    :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-                ]
-            )
-            .reshape(-1, 1, self.tp_v_head_num_, self.head_dim_)
-            .view(torch.float8_e4m3fn)
-        )
-        o = flash_attn_with_kvcache(
-            q=q.view(-1, self.tp_q_head_num_, self.head_dim_),
-            k_cache=cache_k,
-            v_cache=cache_v,
-            page_table=infer_state.page_table,
-            cache_seqlens=infer_state.b_seq_len,
-            cu_seqlens_q=infer_state.cu_seqlens_q,
-            cu_seqlens_k_new=infer_state.cu_seqlens_k,
-            max_seqlen_q=infer_state.q_max_seq_len,
-            causal=True,
-            window_size=(-1, -1),
-            softcap=0.0,
-            q_descale=q_scale,
-            k_descale=infer_state.k_descale[self.layer_num_],
-            v_descale=infer_state.v_descale[self.layer_num_],
-            return_softmax_lse=False,
-        )
-        return o
-
     def _get_o(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight
     ) -> torch.Tensor:
@@ -486,453 +205,6 @@ def _tpsp_ffn(
     #     gate_out, up_out = None, None
     #     return ffn2_out
 
-    def _copy_kv_to_mem_cache_normal(self, buffer, mem_index, mem_manager):
-        destindex_copy_kv(buffer, mem_index, mem_manager.kv_buffer[self.layer_num_])
-        return
-
-    def _copy_kv_to_mem_cache_with_calibration(self, buffer, mem_index, mem_manager):
-        destindex_copy_kv(buffer, mem_index, mem_manager.kv_buffer[self.layer_num_])
-        mem_manager.update_calibration_data(buffer, self.layer_num_)
-        return
-
-    def _copy_kv_to_mem_cache_int8kv(self, buffer, mem_index, mem_manager):
-        destindex_copy_quantize_kv(
-            buffer, mem_index, mem_manager.kv_buffer[self.layer_num_], mem_manager.scale_buffer[self.layer_num_]
-        )
-        return
-
-    def _copy_kv_to_mem_cache_fp8kv(self, buffer, mem_index, mem_manager):
-        scales = mem_manager.scales
-        destindex_copy_kv_fp8(
-            buffer,
-            mem_index,
-            scales[self.layer_num_] if scales is not None else None,
-            mem_manager.kv_buffer[self.layer_num_].view(torch.float8_e4m3fn),
-        )
-        return
-
-    def _copy_kv_to_mem_cache_ppl_int8kv(self, buffer, mem_index, mem_manager):
-        from lightllm.models.llama.triton_kernel.ppl_quant_copy_kv import destindex_copy_quantize_kv
-
-        destindex_copy_quantize_kv(
-            buffer, mem_index, mem_manager.kv_buffer[self.layer_num_], mem_manager.scale_buffer[self.layer_num_]
-        )
-        return
-
-    def _copy_kv_to_mem_cache_ppl_int4kv(self, buffer, mem_index, mem_manager):
-        from lightllm.models.llama.triton_kernel.ppl_int4kv_copy_kv import destindex_copy_int4kv
-
-        destindex_copy_int4kv(
-            buffer, mem_index, mem_manager.kv_buffer[self.layer_num_], mem_manager.scale_buffer[self.layer_num_]
-        )
-        return
-
-    def _token_decode_attention_flashinfer_fp8(self, q, infer_state: LlamaFlashInferStateInfo, layer_weight, out=None):
-        batch_size = infer_state.batch_size
-        calcu_shape1 = (batch_size, self.tp_q_head_num_, self.head_dim_)
-
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_].unsqueeze(1)
-        k = kv[:, :, : self.tp_k_head_num_, :].view(torch.float8_e4m3fn)
-        v = kv[:, :, self.tp_k_head_num_ :, :].view(torch.float8_e4m3fn)
-        offline_scales = infer_state.mem_manager.scales_list
-        k_descale = offline_scales[self.layer_num_][0] if offline_scales is not None else None
-        v_descale = offline_scales[self.layer_num_][1] if offline_scales is not None else None
-        infer_state.decode_wrapper.run(
-            q.view(calcu_shape1),
-            (k, v),
-            k_scale=k_descale,
-            v_scale=v_descale,
-            out=o_tensor.view(calcu_shape1),
-        )
-        return o_tensor
-
-    def _token_decode_attention_flashinfer(self, q, infer_state: LlamaFlashInferStateInfo, layer_weight, out=None):
-        batch_size = infer_state.batch_size
-        calcu_shape1 = (batch_size, self.tp_q_head_num_, self.head_dim_)
-
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_].unsqueeze(1)
-        infer_state.decode_wrapper.run(
-            q.view(calcu_shape1),
-            (kv[:, :, : self.tp_k_head_num_, :], kv[:, :, self.tp_k_head_num_ :, :]),
-            out=o_tensor.view(calcu_shape1),
-        )
-        return o_tensor
-
-    def _token_decode_attention_normal(self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None):
-        total_token_num = infer_state.total_token_num
-        batch_size = infer_state.batch_size
-        calcu_shape1 = (batch_size, self.tp_q_head_num_, self.head_dim_)
-
-        att_m_tensor = self.alloc_tensor((self.tp_q_head_num_, total_token_num), torch.float32)
-
-        token_att_fwd(
-            q.view(calcu_shape1),
-            infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :],
-            att_m_tensor,
-            infer_state.req_manager.req_to_token_indexs,
-            infer_state.b_req_idx,
-            infer_state.b_start_loc,
-            infer_state.b_seq_len,
-            infer_state.max_len_in_batch,
-        )
-
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        from lightllm.models.llama.triton_kernel.token_attention_softmax_and_reducev import (
-            token_softmax_reducev_fwd,
-        )
-
-        token_softmax_reducev_fwd(
-            att_m_tensor,
-            infer_state.mem_manager.kv_buffer[self.layer_num_][
-                :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-            ],
-            o_tensor.view(calcu_shape1),
-            infer_state.req_manager.req_to_token_indexs,
-            infer_state.b_req_idx,
-            infer_state.b_start_loc,
-            infer_state.b_seq_len,
-        )
-        return o_tensor
-
-    def _token_decode_gqa_attention_normal(self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None):
-        batch_size = infer_state.batch_size
-        calcu_shape1 = (batch_size, self.tp_q_head_num_, self.head_dim_)
-        # 对 gqa模型进行推理优化的代码
-        from ..triton_kernel.gqa_decode_flashattention_nopad import gqa_decode_attention_fwd
-
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        gqa_decode_attention_fwd(
-            q.view(calcu_shape1),
-            infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :],
-            infer_state.mem_manager.kv_buffer[self.layer_num_][
-                :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-            ],
-            o_tensor.view(calcu_shape1),
-            infer_state.req_manager.req_to_token_indexs,
-            infer_state.b_req_idx,
-            infer_state.b_seq_len,
-        )
-        return o_tensor
-
-    def _token_decode_attention_int8kv(self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None):
-        total_token_num = infer_state.total_token_num
-        batch_size = infer_state.batch_size
-        calcu_shape1 = (batch_size, self.tp_q_head_num_, self.head_dim_)
-        att_m_tensor = self.alloc_tensor((self.tp_q_head_num_, total_token_num), q.dtype)
-        token_att_fwd_int8k(
-            q.view(calcu_shape1),
-            infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :],
-            infer_state.mem_manager.scale_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :],
-            att_m_tensor,
-            infer_state.req_manager.req_to_token_indexs,
-            infer_state.b_req_idx,
-            infer_state.b_start_loc,
-            infer_state.b_seq_len,
-            infer_state.max_len_in_batch,
-        )
-
-        prob = self.alloc_tensor(att_m_tensor.shape, att_m_tensor.dtype)
-        token_softmax_fwd(
-            att_m_tensor, infer_state.b_start_loc, infer_state.b_seq_len, prob, infer_state.max_len_in_batch
-        )
-        att_m_tensor = None
-
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        token_att_fwd2_int8v(
-            prob,
-            infer_state.mem_manager.kv_buffer[self.layer_num_][
-                :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-            ],
-            infer_state.mem_manager.scale_buffer[self.layer_num_][
-                :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-            ],
-            o_tensor.view(calcu_shape1),
-            infer_state.req_manager.req_to_token_indexs,
-            infer_state.b_req_idx,
-            infer_state.b_start_loc,
-            infer_state.b_seq_len,
-            infer_state.max_len_in_batch,
-        )
-        prob = None
-        return o_tensor
-
-    def _token_decode_attention_flashdecoding(self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None):
-        from lightllm.models.llama.triton_kernel.flash_decoding import token_decode_attention_flash_decoding
-
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :]
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ]
-        return token_decode_attention_flash_decoding(
-            q,
-            infer_state,
-            self.tp_q_head_num_,
-            self.head_dim_,
-            cache_k,
-            cache_v,
-            out=out,
-            alloc_tensor_func=self.alloc_tensor,
-        )
-
-    def _token_decode_attention_gqa_flashdecoding(self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None):
-        # 对 gqa 模型进行推理优化的代码
-        from ..triton_kernel.gqa_flash_decoding import gqa_token_decode_attention_flash_decoding
-
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :]
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ]
-        return gqa_token_decode_attention_flash_decoding(
-            q,
-            infer_state,
-            self.tp_q_head_num_,
-            self.head_dim_,
-            cache_k,
-            cache_v,
-            out=out,
-            alloc_tensor_func=self.alloc_tensor,
-        )
-
-    def _token_decode_attention_ppl_int8kv(self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None):
-        batch_size = infer_state.batch_size
-        calcu_shape1 = (batch_size, self.tp_q_head_num_, self.head_dim_)
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-
-        # group_int8kv_decode_attention(at::Tensor o, at::Tensor q, at::Tensor k, at::Tensor k_s,  at::Tensor v,
-        # at::Tensor v_s, at::Tensor b_loc, at::Tensor b_seq_len, int max_len_in_batch)
-        light_ops.group8_int8kv_decode_attention(
-            o_tensor.view(calcu_shape1),
-            q.view(calcu_shape1),
-            infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :],
-            infer_state.mem_manager.scale_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :],
-            infer_state.mem_manager.kv_buffer[self.layer_num_][
-                :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-            ],
-            infer_state.mem_manager.scale_buffer[self.layer_num_][
-                :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-            ],
-            infer_state.req_manager.req_to_token_indexs,
-            infer_state.b_req_idx,
-            infer_state.b_seq_len,
-            infer_state.max_len_in_batch,
-        )
-
-        return o_tensor
-
-    def _token_decode_attention_ppl_fp16(self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None):
-        batch_size = infer_state.batch_size
-        calcu_shape1 = (batch_size, self.tp_q_head_num_, self.head_dim_)
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        from lightllm_ppl_fp16_kernel import fp16_decode_attention
-
-        # group_int8kv_decode_attention(at::Tensor o, at::Tensor q, at::Tensor k, at::Tensor k_s,
-        # at::Tensor v,  at::Tensor v_s, at::Tensor b_loc, at::Tensor b_seq_len, int max_len_in_batch)
-        fp16_decode_attention(
-            o_tensor.view(calcu_shape1),
-            1.0 / (self.head_dim_ ** 0.5),
-            q.view(calcu_shape1),
-            infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :],
-            infer_state.mem_manager.kv_buffer[self.layer_num_][
-                :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-            ],
-            infer_state.req_manager.req_to_token_indexs,
-            infer_state.b_req_idx,
-            infer_state.b_seq_len,
-            infer_state.max_len_in_batch,
-        )
-
-        return o_tensor
-
-    def _token_decode_attention_ppl_fp16_flashdecoding(
-        self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None
-    ):
-        from lightllm.models.llama.triton_kernel.ppl_fp16_flash_decoding import token_decode_attention_flash_decoding
-
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :]
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ]
-        return token_decode_attention_flash_decoding(
-            q,
-            infer_state,
-            self.tp_q_head_num_,
-            self.head_dim_,
-            cache_k,
-            cache_v,
-            out=out,
-            alloc_tensor_func=self.alloc_tensor,
-        )
-
-    def _token_decode_attention_ppl_int8kv_flashdecoding(
-        self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None
-    ):
-        from lightllm.models.llama.triton_kernel.ppl_int8kv_flash_decoding import token_decode_attention_flash_decoding
-
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :]
-        cache_k_scale = infer_state.mem_manager.scale_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :]
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ]
-        cache_v_scale = infer_state.mem_manager.scale_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ]
-        return token_decode_attention_flash_decoding(
-            q,
-            infer_state,
-            self.tp_q_head_num_,
-            self.head_dim_,
-            cache_k,
-            cache_k_scale,
-            cache_v,
-            cache_v_scale,
-            out=out,
-            alloc_tensor_func=self.alloc_tensor,
-        )
-
-    def _token_decode_attention_ppl_int8kv_flashdecoding_diverse(
-        self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None
-    ):
-        from lightllm.models.llama.triton_kernel.ppl_int8kv_flash_decoding_diverse import (
-            token_decode_attention_flash_decoding,
-        )
-
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :]
-        cache_k_scale = infer_state.mem_manager.scale_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :]
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ]
-        cache_v_scale = infer_state.mem_manager.scale_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ]
-        return token_decode_attention_flash_decoding(
-            q,
-            infer_state,
-            self.tp_q_head_num_,
-            self.head_dim_,
-            cache_k,
-            cache_k_scale,
-            cache_v,
-            cache_v_scale,
-            out=out,
-            alloc_tensor_func=self.alloc_tensor,
-        )
-
-    def _token_decode_attention_ppl_int4kv_flashdecoding(
-        self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None
-    ):
-        from lightllm.models.llama.triton_kernel.ppl_int4kv_flash_decoding import token_decode_attention_flash_decoding
-
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :]
-        cache_k_scale = infer_state.mem_manager.scale_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :]
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ]
-        cache_v_scale = infer_state.mem_manager.scale_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ]
-        return token_decode_attention_flash_decoding(
-            q,
-            infer_state,
-            self.tp_q_head_num_,
-            self.head_dim_,
-            cache_k,
-            cache_k_scale,
-            cache_v,
-            cache_v_scale,
-            out=out,
-            alloc_tensor_func=self.alloc_tensor,
-        )
-
-    def _token_decode_attention_gqa_flashdecoding_vsm(
-        self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None
-    ):
-        from lightllm.models.llama.triton_kernel.gqa_flash_decoding_vsm import (
-            gqa_token_decode_attention_flash_decoding_vsm,
-        )
-
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :]
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ]
-        q_shape = (infer_state.batch_size, self.tp_q_head_num_, self.head_dim_)
-        return gqa_token_decode_attention_flash_decoding_vsm(
-            q.view(q_shape),
-            cache_k,
-            cache_v,
-            infer_state,
-            out=out,
-            alloc_tensor_func=self.alloc_tensor,
-        )
-
-    def _token_decode_attention_flashattention(self, q, infer_state: FlashAttentionStateInfo, layer_weight, out=None):
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :].reshape(
-            -1, 1, self.tp_k_head_num_, self.head_dim_
-        )
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ].reshape(-1, 1, self.tp_v_head_num_, self.head_dim_)
-        q = q.reshape(-1, self.tp_q_head_num_, self.head_dim_)
-        k_descale, v_descale = None, None  # disable quantization
-        Lq = q.shape[-1]
-        sm_scale = 1.0 / (Lq ** 0.5)
-        o = flash_attn_with_kvcache(
-            q=q,
-            k_cache=cache_k,
-            v_cache=cache_v,
-            page_table=infer_state.page_table,
-            cache_seqlens=infer_state.b_att_seq_len,
-            cu_seqlens_q=infer_state.cu_seqlens_q,
-            cu_seqlens_k_new=infer_state.cu_seqlens_k,
-            max_seqlen_q=infer_state.max_q_seq_len,
-            softmax_scale=sm_scale,
-            causal=True,
-            window_size=(-1, -1),
-            softcap=0.0,
-            k_descale=k_descale,
-            v_descale=v_descale,
-            return_softmax_lse=False,
-        )
-        return o
-
-    def _token_decode_attention_flashattention_fp8(
-        self, q, infer_state: FlashAttentionStateInfo, layer_weight, out=None
-    ):
-        cache_k = (
-            (infer_state.mem_manager.kv_buffer[self.layer_num_][:, : self.tp_k_head_num_, :])
-            .reshape(-1, 1, self.tp_k_head_num_, self.head_dim_)
-            .view(torch.float8_e4m3fn)
-        )
-        cache_v = (
-            (
-                infer_state.mem_manager.kv_buffer[self.layer_num_][
-                    :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-                ]
-            )
-            .reshape(-1, 1, self.tp_v_head_num_, self.head_dim_)
-            .view(torch.float8_e4m3fn)
-        )
-        q, q_scale = scaled_fp8_quant(q.view(q.shape[0] * self.tp_k_head_num_, -1), use_per_token_if_dynamic=True)
-        o = flash_attn_with_kvcache(
-            q=q.view(-1, self.tp_q_head_num_, self.head_dim_),
-            k_cache=cache_k,
-            v_cache=cache_v,
-            page_table=infer_state.page_table,
-            cache_seqlens=infer_state.b_seq_len,
-            cu_seqlens_q=infer_state.cu_seqlens_q,
-            cu_seqlens_k_new=infer_state.cu_seqlens_k,
-            max_seqlen_q=1,
-            causal=False,
-            window_size=(-1, -1),
-            softcap=0.0,
-            q_descale=q_scale.view(infer_state.batch_size, self.tp_k_head_num_),
-            k_descale=infer_state.k_descale[self.layer_num_],
-            v_descale=infer_state.v_descale[self.layer_num_],
-            return_softmax_lse=False,
-        )
-        return o
-
     def overlap_tpsp_token_forward(
         self,
         input_embdings: torch.Tensor,
diff --git a/lightllm/models/llama/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/llama/layer_weights/pre_and_post_layer_weight.py
index ea59d24dfc..8efa36cf80 100644
--- a/lightllm/models/llama/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/llama/layer_weights/pre_and_post_layer_weight.py
@@ -1,27 +1,30 @@
 from lightllm.common.basemodel import PreAndPostLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, NoTpNormWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, RMSNormWeight
 
 
 class LlamaPreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
 
+        hidden_size = network_config["hidden_size"]
+        vocab_size = network_config["vocab_size"]
         self.wte_weight_ = EmbeddingWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
             weight_name="model.embed_tokens.weight",
             data_type=self.data_type_,
         )
         tie_word_embeddings = self.network_config_.get("tie_word_embeddings", False)
-        if tie_word_embeddings:
-            self.lm_head_weight_: LMHeadWeight = self.wte_weight_
-        else:
-            self.lm_head_weight_ = LMHeadWeight(
-                weight_name="lm_head.weight",
-                data_type=self.data_type_,
-            )
-
-        self.final_norm_weight_ = NoTpNormWeight(
+        self.lm_head_weight_ = LMHeadWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
+            weight_name="lm_head.weight",
+            data_type=self.data_type_,
+            embedding_weight=self.wte_weight_ if tie_word_embeddings else None,
+        )
+        self.final_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="model.norm.weight",
             data_type=self.data_type_,
-            bias_name=None,
         )
         return
diff --git a/lightllm/models/llama/layer_weights/transformer_layer_weight.py b/lightllm/models/llama/layer_weights/transformer_layer_weight.py
index 6b92272ee7..b68903ecd1 100644
--- a/lightllm/models/llama/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/llama/layer_weights/transformer_layer_weight.py
@@ -2,7 +2,7 @@
 import math
 import numpy as np
 from lightllm.common.basemodel import TransformerLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import ROWMMWeight, COLMMWeight, NoTpNormWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import ROWMMWeight, COLMMWeight, RMSNormWeight, KVROWNMMWeight
 
 
 class LlamaTransformerLayerWeight(TransformerLayerWeight):
@@ -11,10 +11,9 @@ def __init__(
         layer_num,
         data_type,
         network_config,
-        mode=[],
         quant_cfg=None,
     ):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _init_weight(self):
@@ -24,11 +23,15 @@ def _init_weight(self):
         self._init_norm()
 
     def _parse_config(self):
-        self.n_embed = self.network_config_["hidden_size"]
         self.n_head = self.network_config_["num_attention_heads"]
+        self.q_head_num_ = self.network_config_["num_attention_heads"]
+        self.k_head_num_ = self.network_config_["num_key_value_heads"]
+        self.v_head_num_ = self.k_head_num_
+        self.o_head_num_ = self.q_head_num_
+        head_dim = self.network_config_["hidden_size"] // self.network_config_["num_attention_heads"]
+        self.head_dim = self.network_config_.get("head_dim", head_dim)
+        self.n_embed = self.network_config_["hidden_size"]
         self.n_inter = self.network_config_["intermediate_size"]
-        self.n_kv_head = self.network_config_["num_key_value_heads"]
-        self.head_dim = self.network_config_.get("head_dim", self.n_embed // self.n_head)
 
     def _init_weight_names(self):
         self._q_weight_name = f"model.layers.{self.layer_num_}.self_attn.q_proj.weight"
@@ -57,55 +60,67 @@ def _init_weight_names(self):
         self._ffn_norm_bias_name = None
 
     def _init_qkv(self):
+        in_dim = self.n_embed
+        q_out_dim = self.q_head_num_ * self.head_dim
         self.q_proj = ROWMMWeight(
+            in_dim=in_dim,
+            out_dims=[q_out_dim],
             weight_names=self._q_weight_name,
             data_type=self.data_type_,
             bias_names=self._q_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="q_proj",
+            quant_method=self.get_quant_method("q_proj"),
         )
-        self.kv_proj = ROWMMWeight(
+        self.kv_proj = KVROWNMMWeight(
+            in_dim=in_dim,
+            kv_head_num=self.k_head_num_,
+            head_dim=self.head_dim,
             weight_names=[self._k_weight_name, self._v_weight_name],
             data_type=self.data_type_,
             bias_names=[self._k_bias_name, self._v_bias_name],
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="kv_proj",
+            quant_method=self.get_quant_method("kv_proj"),
         )
 
     def _init_o(self):
+        in_dim = self.o_head_num_ * self.head_dim
+        out_dim = self.n_embed
         self.o_proj = COLMMWeight(
+            in_dim=in_dim,
+            out_dims=[out_dim],
             weight_names=self._o_weight_name,
             data_type=self.data_type_,
             bias_names=self._o_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="o_proj",
+            quant_method=self.get_quant_method("o_proj"),
         )
 
     def _init_ffn(self):
         self.gate_up_proj = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_inter, self.n_inter],
             weight_names=[self._gate_weight_name, self._up_weight_name],
             data_type=self.data_type_,
             bias_names=[self._gate_bias_name, self._up_bias_name],
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="gate_up_proj",
+            quant_method=self.get_quant_method("gate_up_proj"),
         )
         self.down_proj = COLMMWeight(
+            in_dim=self.n_inter,
+            out_dims=[self.n_embed],
             weight_names=self._down_weight_name,
             data_type=self.data_type_,
             bias_names=self._down_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="down_proj",
+            quant_method=self.get_quant_method("down_proj"),
         )
 
     def _init_norm(self):
-        self.att_norm_weight_ = NoTpNormWeight(
-            self._att_norm_weight_name, self.data_type_, bias_name=self._att_norm_bias_name
+        hidden_size = self.network_config_["hidden_size"]
+        self.att_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
+            weight_name=self._att_norm_weight_name,
+            data_type=self.data_type_,
+            bias_name=self._att_norm_bias_name,
         )
-        self.ffn_norm_weight_ = NoTpNormWeight(
-            self._ffn_norm_weight_name, self.data_type_, bias_name=self._ffn_norm_bias_name
+        self.ffn_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
+            weight_name=self._ffn_norm_weight_name,
+            data_type=self.data_type_,
+            bias_name=self._ffn_norm_bias_name,
         )
diff --git a/lightllm/models/llama/model.py b/lightllm/models/llama/model.py
index 95465a9e6c..c104ebccc9 100644
--- a/lightllm/models/llama/model.py
+++ b/lightllm/models/llama/model.py
@@ -9,39 +9,15 @@
 from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
 from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.models.llama.flashattention_infer_struct import FlashAttentionStateInfo
-from lightllm.models.llama.flashinfer_struct import LlamaFlashInferStateInfo
 from lightllm.common.basemodel import TpPartBaseModel
 from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class
 from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.envs_utils import get_env_start_args
-from lightllm.utils.dist_utils import get_dp_world_size, get_current_device_id
 
 logger = init_logger(__name__)
 
 
-class LlamaFlashInferStateExtraInfo:
-    def __init__(self, model):
-        tp_world_size = get_dp_world_size()
-        self.tp_q_head_num = model.config["num_attention_heads"] // tp_world_size
-        self.tp_kv_head_num = max(model.config["num_key_value_heads"] // tp_world_size, 1)
-        head_dim = model.config["hidden_size"] // model.config["num_attention_heads"]
-        self.head_dim = model.config.get("head_dim", head_dim)
-        self.workspace_buffer = torch.empty(512 * 1024 * 1024, dtype=torch.int8, device=get_current_device_id())
-        self.max_seq_length = model.max_seq_length
-        self.kv_indices_buffer = [
-            torch.empty(
-                model.graph_max_batch_size * self.max_seq_length, dtype=torch.int32, device=get_current_device_id()
-            ),
-            torch.empty(
-                model.graph_max_batch_size * self.max_seq_length, dtype=torch.int32, device=get_current_device_id()
-            ),
-        ]
-        self.q_data_type = model.data_type
-        self.kv_data_type = torch.float8_e4m3fn if "offline_calibration_fp8kv" in model.mode else model.data_type
-
-
 @ModelRegistry("llama")
 class LlamaTpPartModel(TpPartBaseModel):
     # weight class
@@ -57,9 +33,6 @@ class LlamaTpPartModel(TpPartBaseModel):
     infer_state_class = LlamaInferStateInfo
 
     def __init__(self, kvargs):
-        self.enable_flashinfer = (
-            get_env_start_args().enable_flashinfer_prefill or get_env_start_args().enable_flashinfer_decode
-        )
         super().__init__(kvargs)
         return
 
@@ -94,13 +67,6 @@ def _init_mem_manager(self):
         )
         return
 
-    def _init_inferstate_cls(self):
-        if get_env_start_args().enable_fa3:
-            self.infer_state_class = FlashAttentionStateInfo
-        elif self.enable_flashinfer:
-            self.infer_state_class = LlamaFlashInferStateInfo
-            self.flashinfer_extra_state = LlamaFlashInferStateExtraInfo(self)
-
     def _init_custom(self):
         """
         模型特殊的一些初始化
diff --git a/lightllm/models/llama/triton_kernel/flash_decoding.py b/lightllm/models/llama/triton_kernel/flash_decoding.py
deleted file mode 100644
index e47e308864..0000000000
--- a/lightllm/models/llama/triton_kernel/flash_decoding.py
+++ /dev/null
@@ -1,37 +0,0 @@
-import torch
-
-
-def token_decode_attention_flash_decoding(
-    q, infer_state, q_head_num, head_dim, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty
-):
-    BLOCK_SEQ = 256
-    batch_size = infer_state.batch_size
-    max_len_in_batch = infer_state.max_len_in_batch
-    calcu_shape1 = (batch_size, q_head_num, head_dim)
-
-    from .flash_decoding_stage1 import flash_decode_stage1
-    from .flash_decoding_stage2 import flash_decode_stage2
-
-    o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out
-
-    mid_o = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1, head_dim], dtype=torch.float32, device="cuda"
-    )
-    mid_o_logexpsum = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1], dtype=torch.float32, device="cuda"
-    )
-
-    flash_decode_stage1(
-        q.view(calcu_shape1),
-        cache_k,
-        cache_v,
-        infer_state.req_manager.req_to_token_indexs,
-        infer_state.b_req_idx,
-        infer_state.b_seq_len,
-        infer_state.max_len_in_batch,
-        mid_o,
-        mid_o_logexpsum,
-        BLOCK_SEQ,
-    )
-    flash_decode_stage2(mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ)
-    return o_tensor
diff --git a/lightllm/models/llama/triton_kernel/flash_decoding_stage1.py b/lightllm/models/llama/triton_kernel/flash_decoding_stage1.py
deleted file mode 100644
index 86a3af103d..0000000000
--- a/lightllm/models/llama/triton_kernel/flash_decoding_stage1.py
+++ /dev/null
@@ -1,106 +0,0 @@
-import torch
-import triton
-import triton.language as tl
-
-@triton.jit
-def _fwd_kernel_flash_decode_stage1(
-    Q, K, V, sm_scale, Req_to_tokens, B_req_idx, B_Seqlen,
-    Mid_O, # [batch, head, seq_block_num, head_dim]
-    Mid_O_LogExpSum, #[batch, head, seq_block_num]
-    stride_req_to_tokens_b, stride_req_to_tokens_s,
-    stride_qbs, stride_qh, stride_qd,
-    stride_kbs, stride_kh, stride_kd,
-    stride_vbs, stride_vh, stride_vd,
-    stride_mid_ob, stride_mid_oh, stride_mid_os, stride_mid_od,
-    stride_mid_o_eb, stride_mid_o_eh, stride_mid_o_es,
-    gqa_group_size,
-    BLOCK_SEQ: tl.constexpr, 
-    BLOCK_DMODEL: tl.constexpr,
-    BLOCK_N: tl.constexpr
-):
-    cur_batch = tl.program_id(0)
-    cur_head = tl.program_id(1)
-    seq_start_block = tl.program_id(2)
-    cur_kv_head = cur_head // gqa_group_size
-
-    offs_d = tl.arange(0, BLOCK_DMODEL)
-    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
-    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
-    cur_batch_start_index = seq_start_block * BLOCK_SEQ
-    cur_batch_end_index = tl.minimum(cur_batch_seq_len, cur_batch_start_index + BLOCK_SEQ)
-
-    off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d
-    
-    block_n_size = tl.where(cur_batch_end_index - cur_batch_start_index <= 0, 0, cur_batch_end_index - cur_batch_start_index + BLOCK_N - 1) // BLOCK_N
-    
-    offs_n = cur_batch_start_index + tl.arange(0, BLOCK_N)
-    
-    q = tl.load(Q + off_q)
-
-    sum_exp = 0.0
-    max_logic = -float("inf")
-    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
-
-    for start_n in range(0, block_n_size, 1):
-        offs_n_new = start_n * BLOCK_N + offs_n
-        k_loc = tl.load(Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx +  offs_n_new, mask=offs_n_new < cur_batch_end_index, other=0)
-        k_loc = k_loc.to(tl.int64)
-        off_k = k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :]
-        k = tl.load(K + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0)
-        att_value = tl.sum(q[None, :] * k, 1)
-        att_value *= sm_scale
-        att_value = tl.where(offs_n_new < cur_batch_end_index, att_value, float("-inf"))
-        v = tl.load(V + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0)
-        
-        cur_max_logic = tl.max(att_value, axis=0)
-        new_max_logic = tl.maximum(cur_max_logic, max_logic)
-
-        exp_logic = tl.exp(att_value - new_max_logic)
-        logic_scale = tl.exp(max_logic - new_max_logic)
-        acc *= logic_scale
-        acc += tl.sum(exp_logic[:, None] * v, axis=0)
-
-        sum_exp = sum_exp * logic_scale + tl.sum(exp_logic, axis=0)
-        max_logic = new_max_logic
-    
-    need_store = tl.where(block_n_size == 0, 0, 1)
-    for _ in range(0, need_store, 1):
-        off_mid_o = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + seq_start_block * stride_mid_os + offs_d
-        off_mid_o_logexpsum = cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh + seq_start_block
-        tl.store(Mid_O + off_mid_o, acc / sum_exp)
-        tl.store(Mid_O_LogExpSum + off_mid_o_logexpsum, max_logic + tl.log(sum_exp))
-    return
-
-
-@torch.no_grad()
-def flash_decode_stage1(q, k, v, Req_to_tokens, B_req_idx, B_Seqlen, max_len_in_batch, mid_out, mid_out_logsumexp, block_seq):
-    BLOCK_SEQ = block_seq
-    BLOCK_N = 16
-    assert BLOCK_SEQ % BLOCK_N == 0
-    # shape constraints
-    Lq, Lk = q.shape[-1], k.shape[-1]
-    assert Lq == Lk
-    assert Lk in {16, 32, 64, 128}
-    sm_scale = 1.0 / (Lk ** 0.5)
-    batch, head_num = B_req_idx.shape[0], q.shape[1]
-    grid = (batch, head_num, triton.cdiv(max_len_in_batch, BLOCK_SEQ))
-    gqa_group_size = q.shape[1] // k.shape[1]
-    
-    _fwd_kernel_flash_decode_stage1[grid](
-        q, k, v, sm_scale, Req_to_tokens, B_req_idx, B_Seqlen,
-        mid_out,
-        mid_out_logsumexp,
-        Req_to_tokens.stride(0), Req_to_tokens.stride(1),
-        q.stride(0), q.stride(1), q.stride(2),
-        k.stride(0), k.stride(1), k.stride(2),
-        v.stride(0), v.stride(1), v.stride(2),
-        mid_out.stride(0), mid_out.stride(1), mid_out.stride(2), mid_out.stride(3),
-        mid_out_logsumexp.stride(0), mid_out_logsumexp.stride(1), mid_out_logsumexp.stride(2),
-        gqa_group_size,
-        BLOCK_SEQ=BLOCK_SEQ,
-        BLOCK_DMODEL=Lk,
-        BLOCK_N=BLOCK_N,
-        num_warps=1,
-        num_stages=2,
-    )
-    return
\ No newline at end of file
diff --git a/lightllm/models/llama/triton_kernel/flash_decoding_stage2.py b/lightllm/models/llama/triton_kernel/flash_decoding_stage2.py
deleted file mode 100644
index 81227f967b..0000000000
--- a/lightllm/models/llama/triton_kernel/flash_decoding_stage2.py
+++ /dev/null
@@ -1,64 +0,0 @@
-import torch
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def _fwd_kernel_flash_decode_stage2(
-    B_Seqlen,
-    Mid_O, # [batch, head, seq_block_num, head_dim]
-    Mid_O_LogExpSum, #[batch, head, seq_block_num]
-    O, #[batch, head, head_dim]
-    stride_mid_ob, stride_mid_oh, stride_mid_os, stride_mid_od,
-    stride_mid_o_eb, stride_mid_o_eh, stride_mid_o_es,
-    stride_obs, stride_oh, stride_od,
-    BLOCK_SEQ: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr):
-    cur_batch = tl.program_id(0)
-    cur_head = tl.program_id(1)
-
-    offs_d = tl.arange(0, BLOCK_DMODEL)
-    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
-
-    block_n_size = tl.where(cur_batch_seq_len <= 0, 0, cur_batch_seq_len + BLOCK_SEQ - 1) // BLOCK_SEQ
-
-    sum_exp = 0.0
-    max_logic = -float("inf")
-    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
-
-    offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d
-    offs_logic = cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh
-    for block_seq_n in range(0, block_n_size, 1):
-        tv = tl.load(Mid_O + offs_v + block_seq_n * stride_mid_os)
-        tlogic = tl.load(Mid_O_LogExpSum + offs_logic + block_seq_n)
-        new_max_logic = tl.maximum(tlogic, max_logic)
-        
-        old_scale = tl.exp(max_logic - new_max_logic)
-        acc *= old_scale
-        exp_logic = tl.exp(tlogic - new_max_logic)
-        acc += exp_logic * tv
-        sum_exp = sum_exp * old_scale + exp_logic
-        max_logic = new_max_logic
-    
-    tl.store(O + cur_batch * stride_obs + cur_head * stride_oh + offs_d, acc / sum_exp)
-    return
-
-
-@torch.no_grad()
-def flash_decode_stage2(mid_out, mid_out_logexpsum, B_Seqlen, O, block_seq):
-    Lk = mid_out.shape[-1]
-    assert Lk in {16, 32, 64, 128}
-    batch, head_num = mid_out.shape[0], mid_out.shape[1]
-    grid = (batch, head_num)
-    
-    _fwd_kernel_flash_decode_stage2[grid](
-        B_Seqlen, mid_out, mid_out_logexpsum, O,
-        mid_out.stride(0), mid_out.stride(1), mid_out.stride(2), mid_out.stride(3),
-        mid_out_logexpsum.stride(0), mid_out_logexpsum.stride(1), mid_out_logexpsum.stride(2),
-        O.stride(0), O.stride(1), O.stride(2),
-        BLOCK_SEQ=block_seq,
-        BLOCK_DMODEL=Lk,
-        num_warps=4,
-        num_stages=2,
-    )
-    return
\ No newline at end of file
diff --git a/lightllm/models/llama/triton_kernel/ppl_int4kv_copy_kv.py b/lightllm/models/llama/triton_kernel/ppl_int4kv_copy_kv.py
deleted file mode 100644
index 7ba0f3b31b..0000000000
--- a/lightllm/models/llama/triton_kernel/ppl_int4kv_copy_kv.py
+++ /dev/null
@@ -1,138 +0,0 @@
-import torch
-
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def _fwd_kernel_destindex_copy_quantize_int4_kv(
-    K,
-    Dest_loc,
-    Out,
-    Out_scale,
-    stride_k_bs,
-    stride_k_h,
-    stride_k_g,
-    stride_k_d,
-    stride_o_bs,
-    stride_o_h,
-    stride_o_g,
-    stride_o_d,
-    stride_os_bs,
-    stride_os_h,
-    stride_os_g,
-    group_size,
-    BLOCK_GROUP_NUM: tl.constexpr,
-    BLOCK_GROUP_DIM: tl.constexpr,
-):
-    cur_index = tl.program_id(0)
-    cur_head = tl.program_id(1)
-
-    offs_g = tl.arange(0, BLOCK_GROUP_NUM)
-    offs_d = tl.arange(0, BLOCK_GROUP_DIM // 2)
-
-    dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
-
-    src_data_0 = tl.load(
-        K + cur_index * stride_k_bs + cur_head * stride_k_h + offs_g[:, None] * stride_k_g + offs_d[None, :] * 2,
-        mask=offs_g[:, None] < group_size,
-        other=0.0,
-    )
-    src_data_1 = tl.load(
-        K + cur_index * stride_k_bs + cur_head * stride_k_h + offs_g[:, None] * stride_k_g + offs_d[None, :] * 2 + 1,
-        mask=offs_g[:, None] < group_size,
-        other=0.0,
-    )
-
-    abs_data_0 = tl.abs(src_data_0)
-    abs_data_1 = tl.abs(src_data_1)
-
-    data_scale = (tl.maximum(tl.max(abs_data_0, axis=1), tl.max(abs_data_1, axis=1)) / 7.0).to(Out_scale.dtype.element_ty)
-    q_src_data_0 = (src_data_0 / data_scale[:, None]).to(tl.int8)
-    q_src_data_0 = tl.where(q_src_data_0 > 7, 7, q_src_data_0)
-    q_src_data_0 = tl.where(q_src_data_0 < -7, -7, q_src_data_0)
-
-    q_src_data_1 = (src_data_1 / data_scale[:, None]).to(tl.int8)
-    q_src_data_1 = tl.where(q_src_data_1 > 7, 7, q_src_data_1)
-    q_src_data_1 = tl.where(q_src_data_1 < -7, -7, q_src_data_1)
-
-    low_4 = ((q_src_data_0 & 0x80) >> 4) | (q_src_data_0 & 0xF)
-    high_4 = (((q_src_data_1 & 0x80) >> 4) | (q_src_data_1 & 0xF)) << 4
-
-    # tl.device_print(low_4)
-    # tl.device_print(high_4)
-
-    out_data = low_4 | high_4
-
-    o_ptrs = Out + dest_index * stride_o_bs + cur_head * stride_o_h + offs_g[:, None] * stride_o_g + offs_d[None, :]
-    os_ptrs = Out_scale + dest_index * stride_os_bs + cur_head * stride_os_h + offs_g
-    tl.store(o_ptrs, out_data, mask=offs_g[:, None] < group_size)
-    tl.store(os_ptrs, data_scale, mask=offs_g < group_size)
-    return
-
-
-@torch.no_grad()
-def destindex_copy_int4kv(K, DestLoc, Out, Out_scale):
-    # seq_len = DestLoc.shape[0]
-    # head_num = K.shape[1]
-    head_dim = K.shape[2]
-    quant_group_dim = 8
-
-    assert head_dim % quant_group_dim == 0, "error head dim, can not been supported to copy quant kv"
-    # grid = (seq_len, head_num)
-    # num_warps = 1
-
-    group_size = head_dim // quant_group_dim
-    group_dim = quant_group_dim
-
-    K = K.view((K.shape[0], K.shape[1], group_size, group_dim))
-    Out = Out.view(
-        Out.shape[0], Out.shape[1], group_size, group_dim // 2
-    )  # OUt 是 int8 类型， 两个int4组一个int8，所以 group_dim // 2
-
-    from lightllm_ppl_int4kv_flashdecoding_kernel import group8_copy_int4_kv
-
-    group8_copy_int4_kv(Out, Out_scale, K, DestLoc, 4)
-
-    # _fwd_kernel_destindex_copy_quantize_int4_kv[grid](
-    #     K,
-    #     DestLoc,
-    #     Out,
-    #     Out_scale,
-    #     K.stride(0),
-    #     K.stride(1),
-    #     K.stride(2),
-    #     K.stride(3),
-    #     Out.stride(0),
-    #     Out.stride(1),
-    #     Out.stride(2),
-    #     Out.stride(3),
-    #     Out_scale.stride(0),
-    #     Out_scale.stride(1),
-    #     Out_scale.stride(2),
-    #     group_size,
-    #     BLOCK_GROUP_NUM=triton.next_power_of_2(group_size),
-    #     BLOCK_GROUP_DIM=group_dim,
-    #     num_warps=num_warps,
-    #     num_stages=1,
-    # )
-    return
-
-
-def test2():
-    import time
-
-    src = torch.randn((1, 1, 8), dtype=torch.float16).cuda()
-    src[0, 0, :] = torch.tensor([1, -2, 2, 0, 4, 5, 6, 7]).cuda()
-    dest_loc = torch.arange(0, 1, dtype=torch.int32).cuda()
-    value_dest = torch.randn((1, 1, 4), dtype=torch.float16).cuda().to(torch.int8)
-    scale_dest = torch.randn((1, 1, 1), dtype=torch.float16).cuda()
-
-    destindex_copy_int4kv(src, dest_loc, value_dest, scale_dest)
-
-    print(value_dest)
-    print(scale_dest)
-
-
-if __name__ == "__main__":
-    test2()
diff --git a/lightllm/models/llama/triton_kernel/ppl_int4kv_flash_decoding.py b/lightllm/models/llama/triton_kernel/ppl_int4kv_flash_decoding.py
deleted file mode 100644
index 1e324bcc0b..0000000000
--- a/lightllm/models/llama/triton_kernel/ppl_int4kv_flash_decoding.py
+++ /dev/null
@@ -1,50 +0,0 @@
-import torch
-
-
-def token_decode_attention_flash_decoding(
-    q,
-    infer_state,
-    q_head_num,
-    head_dim,
-    cache_k,
-    cache_k_scale,
-    cache_v,
-    cache_v_scale,
-    out=None,
-    alloc_tensor_func=torch.empty,
-):
-    BLOCK_SEQ = 256
-    batch_size = infer_state.batch_size
-    max_len_in_batch = infer_state.max_len_in_batch
-    calcu_shape1 = (batch_size, q_head_num, head_dim)
-
-    from lightllm_ppl_int4kv_flashdecoding_kernel import group8_int4kv_flashdecoding_stage1
-    from .flash_decoding_stage2 import flash_decode_stage2
-
-    o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out
-
-    mid_o = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1, head_dim], dtype=torch.float16, device="cuda"
-    )
-    mid_o_logexpsum = alloc_tensor_func(
-        [batch_size, q_head_num, max_len_in_batch // BLOCK_SEQ + 1], dtype=torch.float16, device="cuda"
-    )
-
-    group8_int4kv_flashdecoding_stage1(
-        BLOCK_SEQ,
-        mid_o,
-        mid_o_logexpsum,
-        1.0 / (head_dim ** 0.5),
-        q.view(calcu_shape1),
-        cache_k,
-        cache_k_scale,
-        cache_v,
-        cache_v_scale,
-        infer_state.req_manager.req_to_token_indexs,
-        infer_state.b_req_idx,
-        infer_state.b_seq_len,
-        infer_state.max_len_in_batch,
-    )
-
-    flash_decode_stage2(mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ)
-    return o_tensor
diff --git a/lightllm/models/llama/triton_kernel/ppl_quant_copy_kv.py b/lightllm/models/llama/triton_kernel/ppl_quant_copy_kv.py
deleted file mode 100644
index 3d9a490f47..0000000000
--- a/lightllm/models/llama/triton_kernel/ppl_quant_copy_kv.py
+++ /dev/null
@@ -1,294 +0,0 @@
-import torch
-
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def _fwd_kernel_destindex_copy_quantize_kv(
-    K,
-    Dest_loc,
-    Out,
-    Out_scale,
-    stride_k_bs,
-    stride_k_h,
-    stride_k_g,
-    stride_k_d,
-    stride_o_bs,
-    stride_o_h,
-    stride_o_g,
-    stride_o_d,
-    stride_os_bs,
-    stride_os_h,
-    stride_os_g,
-    group_size,
-    BLOCK_GROUP_NUM: tl.constexpr,
-    BLOCK_GROUP_DIM: tl.constexpr,
-):
-    cur_index = tl.program_id(0)
-    cur_head = tl.program_id(1)
-
-    offs_g = tl.arange(0, BLOCK_GROUP_NUM)
-    offs_d = tl.arange(0, BLOCK_GROUP_DIM)
-
-    dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
-
-    src_data = tl.load(
-        K + cur_index * stride_k_bs + cur_head * stride_k_h + offs_g[:, None] * stride_k_g + offs_d[None, :],
-        mask=offs_g[:, None] < group_size,
-        other=0.0,
-    )
-    abs_data = tl.abs(src_data)
-    data_scale = (tl.max(abs_data, axis=1) / 127.0).to(Out_scale.dtype.element_ty)
-    q_src_data = (src_data / data_scale[:, None]).to(tl.int8)
-
-    o_ptrs = Out + dest_index * stride_o_bs + cur_head * stride_o_h + offs_g[:, None] * stride_o_g + offs_d[None, :]
-    os_ptrs = Out_scale + dest_index * stride_os_bs + cur_head * stride_os_h + offs_g
-    tl.store(o_ptrs, q_src_data, mask=offs_g[:, None] < group_size)
-    tl.store(os_ptrs, data_scale, mask=offs_g < group_size)
-    return
-
-
-@torch.no_grad()
-def destindex_copy_quantize_kv(K, DestLoc, Out, Out_scale):
-    seq_len = DestLoc.shape[0]
-    head_num = K.shape[1]
-    head_dim = K.shape[2]
-    quant_group_dim = 8
-
-    assert head_dim % quant_group_dim == 0, "error head dim, can not been supported to copy quant kv"
-    grid = (seq_len, head_num)
-    num_warps = 1
-
-    group_size = head_dim // quant_group_dim
-    group_dim = quant_group_dim
-
-    K = K.view((K.shape[0], K.shape[1], group_size, group_dim))
-    Out = Out.view(Out.shape[0], Out.shape[1], group_size, group_dim)
-
-    _fwd_kernel_destindex_copy_quantize_kv[grid](
-        K,
-        DestLoc,
-        Out,
-        Out_scale,
-        K.stride(0),
-        K.stride(1),
-        K.stride(2),
-        K.stride(3),
-        Out.stride(0),
-        Out.stride(1),
-        Out.stride(2),
-        Out.stride(3),
-        Out_scale.stride(0),
-        Out_scale.stride(1),
-        Out_scale.stride(2),
-        group_size,
-        BLOCK_GROUP_NUM=triton.next_power_of_2(group_size),
-        BLOCK_GROUP_DIM=group_dim,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
-
-
-@triton.jit
-def _fwd_kernel_destindex_copy_dequantize_kv(
-    mem_kv_buffer,
-    mem_kv_scale,
-    req_to_token_indexs,
-    b_seq_len,
-    b_req_idx,
-    Out,
-    stride_kv_b,
-    stride_kv_h,
-    stride_kv_g,
-    stride_kv_d,
-    stride_o_bh,
-    stride_o_l,
-    stride_o_g,
-    stride_o_d,
-    stride_s_b,
-    stride_s_h,
-    stride_s_g,
-    stride_req_to_tokens_b,
-    stride_req_to_tokens_s,
-    group_size,
-    head_num: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-    BLOCK_GROUP_NUM: tl.constexpr,
-    BLOCK_GROUP_DIM: tl.constexpr,
-):
-    cur_group = tl.program_id(0)
-    start_m = tl.program_id(1)
-    cur_bh = tl.program_id(2)
-    cur_batch = cur_bh // head_num
-    cur_head = cur_bh % head_num
-
-    block_start_loc = BLOCK_SIZE * start_m
-    cur_batch_req_idx = tl.load(b_req_idx + cur_batch)
-    cur_seq_len = tl.load(b_seq_len + cur_batch)
-
-    # initialize offsets
-    offs_kv_loc = block_start_loc + tl.arange(0, BLOCK_SIZE)
-
-    # offs_g = tl.arange(0, BLOCK_GROUP_NUM)
-    offs_d = tl.arange(0, BLOCK_GROUP_DIM)
-
-    kv_loc = tl.load(
-        req_to_token_indexs + cur_batch_req_idx * stride_req_to_tokens_b + offs_kv_loc, mask=offs_kv_loc < cur_seq_len
-    ).to(tl.int64)
-    offs_kv = kv_loc[:, None] * stride_kv_b + cur_head * stride_kv_h + cur_group * stride_kv_g + offs_d[None, :]
-
-    src_data = tl.load(
-        mem_kv_buffer + offs_kv,
-        mask=offs_kv_loc[:, None] < cur_seq_len,
-        other=0.0,
-    ).to(Out.dtype.element_ty)
-
-    s_ptrs = mem_kv_scale + kv_loc * stride_s_b + cur_head * stride_s_h + cur_group * stride_s_g
-    data_scale = tl.load(
-        s_ptrs,
-        mask=offs_kv_loc < cur_seq_len,
-    )
-
-    out_data = src_data * data_scale[:, None]
-    o_ptrs = Out + cur_bh * stride_o_bh + offs_kv_loc[:, None] * stride_o_l + cur_group * stride_o_g + offs_d[None, :]
-    tl.store(o_ptrs, out_data, mask=offs_kv_loc[:, None] < cur_seq_len)
-    return
-
-
-@torch.no_grad()
-def destindex_copy_dequantize_kv(
-    mem_kv_buffer, mem_kv_scale, req_to_token_indexs, b_seq_len, b_req_idx, max_len_in_batch, Out
-):
-    batch_size = b_seq_len.shape[0]
-    head_num = mem_kv_buffer.shape[1]
-    head_dim = mem_kv_buffer.shape[2]
-    quant_group_dim = 8
-    BLOCK_SIZE = 128
-    group_size = head_dim // quant_group_dim
-    group_dim = quant_group_dim
-    assert head_dim % quant_group_dim == 0, "error head dim, can not been supported to copy quant kv"
-    grid = (group_size, triton.cdiv(max_len_in_batch, BLOCK_SIZE), batch_size * head_num)
-    num_warps = 1
-    mem_kv_buffer = mem_kv_buffer.view((mem_kv_buffer.shape[0], mem_kv_buffer.shape[1], group_size, group_dim))
-    mem_kv_scale = mem_kv_scale.view((mem_kv_buffer.shape[0], mem_kv_buffer.shape[1], -1))
-    Out = Out.view(Out.shape[0] * Out.shape[1], -1, group_size, group_dim)
-
-    _fwd_kernel_destindex_copy_dequantize_kv[grid](
-        mem_kv_buffer,
-        mem_kv_scale,
-        req_to_token_indexs,
-        b_seq_len,
-        b_req_idx,
-        Out,
-        mem_kv_buffer.stride(0),
-        mem_kv_buffer.stride(1),
-        mem_kv_buffer.stride(2),
-        mem_kv_buffer.stride(3),
-        Out.stride(0),
-        Out.stride(1),
-        Out.stride(2),
-        Out.stride(3),
-        mem_kv_scale.stride(0),
-        mem_kv_scale.stride(1),
-        mem_kv_scale.stride(2),
-        req_to_token_indexs.stride(0),
-        req_to_token_indexs.stride(1),
-        group_size,
-        head_num=head_num,
-        BLOCK_SIZE=BLOCK_SIZE,
-        BLOCK_GROUP_NUM=triton.next_power_of_2(group_size),
-        BLOCK_GROUP_DIM=group_dim,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
-
-
-def test2():
-    import time
-
-    B, N_CTX, H, D = 1, 3, 12, 128
-    src = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda()
-    dest_loc = torch.arange(0, B * N_CTX, dtype=torch.int32).cuda()
-    value_dest = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda().to(torch.int8)
-    scale_dest = torch.randn((B * N_CTX, H, D // 8), dtype=torch.float16).cuda()
-
-    for _ in range(10):
-        destindex_copy_quantize_kv(src, dest_loc, value_dest, scale_dest)
-    torch.cuda.synchronize()
-    t1 = time.time()
-    for _ in range(1000):
-        destindex_copy_quantize_kv(src, dest_loc, value_dest, scale_dest)
-    torch.cuda.synchronize()
-    t2 = time.time()
-
-    print("Time cost ", t2 - t1)
-    value_dest = value_dest.view((B * N_CTX, H, D // 8, 8))
-    scale_dest = scale_dest.view((B * N_CTX, H, D // 8, 1))
-    print("max ", torch.max(torch.abs((value_dest * scale_dest).view(B * N_CTX, H, D) - src)))
-    print("mean ", torch.mean(torch.abs((value_dest * scale_dest).view(B * N_CTX, H, D) - src)))
-    cos = torch.nn.CosineSimilarity(0)
-    print("cos ", cos(src.flatten().to(torch.float32), (value_dest * scale_dest).flatten().to(torch.float32)))
-
-
-def torch_dequant(kv, kv_scale, o, b_req_idx, b_seq_len, req_to_token_indexs):
-
-    batch = b_req_idx.shape[0]
-    for i in range(batch):
-        req_idx = b_req_idx[i]
-        seq_len = b_seq_len[i]
-        print(seq_len, b_seq_len)
-        kv_loc = req_to_token_indexs[req_idx, :seq_len]
-        head_num = kv.shape[1]
-        cur_kv = kv[kv_loc, :, :].reshape(seq_len, head_num, -1, 8).to(o.dtype)
-        cur_scale = kv_scale[kv_loc, :, :].reshape(seq_len, head_num, -1, 1)
-        out = cur_kv * cur_scale
-        o[i, :seq_len, :, :] = out.reshape(out.shape[0], out.shape[1], -1)
-
-
-def test3():
-    import time
-    import numpy as np
-
-    Z, H, N_CTX, D_HEAD = 1, 16, 3, 128
-    dtype = torch.bfloat16
-    kv = torch.empty((Z * N_CTX + 100, 2 * H, D_HEAD), dtype=torch.int8, device="cuda")
-    kv_scale = torch.randn((Z * N_CTX + 100, 2 * H, D_HEAD // 8), dtype=dtype, device="cuda")
-    out = torch.empty((Z, 2 * H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.4, std=0.2)
-    torch_out = torch.empty((Z, N_CTX, 2 * H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.3, std=0.2)
-    req_to_token_indexs = torch.empty((1000, N_CTX + 7000), dtype=torch.int32, device="cuda")
-    max_input_len = N_CTX
-    b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda")
-    b_req_idx = torch.ones((Z,), dtype=torch.int32, device="cuda")
-    for i in range(Z):
-        seq_len = N_CTX - i * 100
-        b_seq_len[i] = seq_len
-        b_req_idx[i] = i
-        req_to_token_indexs[i][:seq_len] = (
-            torch.tensor(np.arange(seq_len), dtype=torch.int32).cuda() + b_seq_len[0:i].sum()
-        )
-    print(b_seq_len)
-    destindex_copy_dequantize_kv(kv, kv_scale, req_to_token_indexs, b_seq_len, b_req_idx, max_input_len, out)
-    torch_dequant(kv, kv_scale, torch_out, b_req_idx, b_seq_len, req_to_token_indexs)
-    torch.cuda.synchronize()
-    t1 = time.time()
-    for _ in range(1000):
-        destindex_copy_dequantize_kv(kv, kv_scale, req_to_token_indexs, b_seq_len, b_req_idx, max_input_len, out)
-    torch.cuda.synchronize()
-    t2 = time.time()
-    print((t2 - t1))
-    torch_out = torch_out.transpose(1, 2)
-    for i in range(Z):
-        print("max ", torch.max(torch.abs(torch_out - out)[i][:, : b_seq_len[i]]))
-        print("mean ", torch.mean(torch.abs(torch_out - out)[i][:, : b_seq_len[i]]))
-        assert torch.allclose(torch_out[i][:, : b_seq_len[i]], out[i][:, : b_seq_len[i]], atol=1e-2, rtol=0)
-    # print("max ", torch.max(torch.abs((value_dest * scale_dest).view(B * N_CTX, H, D) - src)))
-    # print("mean ", torch.mean(torch.abs((value_dest * scale_dest).view(B * N_CTX, H, D) - src)))
-    # cos = torch.nn.CosineSimilarity(0)
-    # print("cos ", cos(src.flatten().to(torch.float32), (value_dest * scale_dest).flatten().to(torch.float32)))
-
-
-if __name__ == "__main__":
-    test3()
diff --git a/lightllm/models/llama/triton_kernel/token_attention_nopad_att1.py b/lightllm/models/llama/triton_kernel/token_attention_nopad_att1.py
index eb5af6fecd..45de83e989 100644
--- a/lightllm/models/llama/triton_kernel/token_attention_nopad_att1.py
+++ b/lightllm/models/llama/triton_kernel/token_attention_nopad_att1.py
@@ -60,7 +60,8 @@ def _fwd_kernel_token_att1(
         ).to(tl.int64)
         off_k = k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :] * stride_kd
         k = tl.load(K + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0)
-        att_value = tl.sum(q[None, :] * k, 1, dtype=tl.float32)
+        att_value = tl.sum(q[None, :] * k, 1)
+        att_value = att_value.to(tl.float32)
         att_value *= sm_scale
         off_o = cur_head * att_stride_h + (cur_batch_in_all_start_index + offs_n) * att_stride_bs
         tl.store(Att_Out + off_o, att_value, mask=offs_n_new < cur_batch_end_index)
diff --git a/lightllm/models/llama/triton_kernel/token_attention_nopad_reduceV.py b/lightllm/models/llama/triton_kernel/token_attention_nopad_reduceV.py
deleted file mode 100644
index 243a8d1f66..0000000000
--- a/lightllm/models/llama/triton_kernel/token_attention_nopad_reduceV.py
+++ /dev/null
@@ -1,223 +0,0 @@
-import torch
-
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def _fwd_kernel_token_att2(
-    Prob,
-    V,
-    Out,
-    Req_to_tokens,
-    B_req_idx,
-    B_Start_Loc,
-    B_Seqlen,
-    stride_req_to_tokens_b,
-    stride_req_to_tokens_s,
-    stride_ph,
-    stride_pbs,
-    stride_vbs,
-    stride_vh,
-    stride_vd,
-    stride_obs,
-    stride_oh,
-    stride_od,
-    kv_group_num,
-    BLOCK_DMODEL: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-):
-    cur_batch = tl.program_id(0)
-    cur_head = tl.program_id(1)
-
-    cur_kv_head = cur_head // kv_group_num
-
-    offs_n = tl.arange(0, BLOCK_N)
-    offs_d = tl.arange(0, BLOCK_DMODEL)
-    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
-    cur_batch_start_index = 0
-    cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
-    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
-
-    v_loc_off = cur_batch_req_idx * stride_req_to_tokens_b + (cur_batch_start_index + offs_n) * stride_req_to_tokens_s
-    p_offs = cur_head * stride_ph + (cur_batch_in_all_start_index + offs_n) * stride_pbs
-    v_offs = cur_kv_head * stride_vh + offs_d[None, :] * stride_vd
-
-    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
-    for start_n in range(0, cur_batch_seq_len, BLOCK_N):
-        start_n = tl.multiple_of(start_n, BLOCK_N)
-        p_value = tl.load(Prob + p_offs + start_n, mask=(start_n + offs_n) < cur_batch_seq_len, other=0.0)
-        v_loc = tl.load(
-            Req_to_tokens + v_loc_off + start_n * stride_req_to_tokens_s,
-            mask=(start_n + offs_n) < cur_batch_seq_len,
-            other=0.0,
-        ).to(tl.int64)
-        v_value = tl.load(
-            V + v_offs + v_loc[:, None] * stride_vbs, mask=(start_n + offs_n[:, None]) < cur_batch_seq_len, other=0.0
-        )
-        acc += tl.sum(p_value[:, None] * v_value, 0)
-
-    acc = acc.to(Out.dtype.element_ty)
-    off_o = cur_batch * stride_obs + cur_head * stride_oh + offs_d * stride_od
-    out_ptrs = Out + off_o
-    tl.store(out_ptrs, acc)
-    return
-
-
-@torch.no_grad()
-def token_att_fwd2(prob, v, out, Req_to_tokens, B_req_idx, B_Start_Loc, B_Seqlen):
-    BLOCK = 128
-    # BLOCK = 64 # for triton 2.0.0dev
-    batch, head = B_req_idx.shape[0], prob.shape[0]
-    grid = (batch, head)
-    num_warps = 4
-    dim = v.shape[-1]
-
-    kv_group_num = prob.shape[0] // v.shape[1]
-
-    _fwd_kernel_token_att2[grid](
-        prob,
-        v,
-        out,
-        Req_to_tokens,
-        B_req_idx,
-        B_Start_Loc,
-        B_Seqlen,
-        Req_to_tokens.stride(0),
-        Req_to_tokens.stride(1),
-        prob.stride(0),
-        prob.stride(1),
-        v.stride(0),
-        v.stride(1),
-        v.stride(2),
-        out.stride(0),
-        out.stride(1),
-        out.stride(2),
-        kv_group_num=kv_group_num,
-        BLOCK_DMODEL=dim,
-        BLOCK_N=BLOCK,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
-
-
-@triton.jit
-def _fwd_kernel_token_att2_int8v(
-    Prob,
-    V,
-    V_scale,
-    Out,
-    Req_to_tokens,
-    B_req_idx,
-    B_Start_Loc,
-    B_Seqlen,  # B_Start_Loc 保存的是如果连续存储时候的累加输入和
-    stride_req_to_tokens_b,
-    stride_req_to_tokens_s,
-    stride_ph,
-    stride_pbs,
-    stride_vbs,
-    stride_vh,
-    stride_vd,
-    stride_vsbs,
-    stride_vsh,
-    stride_vsd,
-    stride_obs,
-    stride_oh,
-    stride_od,
-    kv_group_num,
-    BLOCK_DMODEL: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-):
-    cur_batch = tl.program_id(0)
-    cur_head = tl.program_id(1)
-
-    cur_kv_head = cur_head // kv_group_num
-
-    offs_n = tl.arange(0, BLOCK_N)
-    offs_d = tl.arange(0, BLOCK_DMODEL)
-    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
-    cur_batch_start_index = 0
-    cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
-    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
-
-    v_loc_off = cur_batch_req_idx * stride_req_to_tokens_b + (cur_batch_start_index + offs_n) * stride_req_to_tokens_s
-    p_offs = cur_head * stride_ph + (cur_batch_in_all_start_index + offs_n) * stride_pbs
-    v_offs = cur_kv_head * stride_vh + offs_d[None, :] * stride_vd
-    vs_offs = cur_kv_head * stride_vsh
-
-    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
-    for start_n in range(0, cur_batch_seq_len, BLOCK_N):
-        start_n = tl.multiple_of(start_n, BLOCK_N)
-        p_value = tl.load(Prob + p_offs + start_n, mask=(start_n + offs_n) < cur_batch_seq_len, other=0.0)
-        v_loc = tl.load(
-            Req_to_tokens + v_loc_off + start_n * stride_req_to_tokens_s,
-            mask=(start_n + offs_n) < cur_batch_seq_len,
-            other=0.0,
-        ).to(tl.int64)
-        v_value = tl.load(
-            V + v_offs + v_loc[:, None] * stride_vbs, mask=(start_n + offs_n[:, None]) < cur_batch_seq_len, other=0.0
-        )
-        vs_value = tl.load(
-            V_scale + vs_offs + v_loc[:, None] * stride_vsbs,
-            mask=(start_n + offs_n[:, None]) < cur_batch_seq_len,
-            other=0.0,
-        )
-        acc += tl.sum(p_value[:, None] * v_value * vs_value, 0)
-
-    acc = acc.to(Out.dtype.element_ty)
-    off_o = cur_batch * stride_obs + cur_head * stride_oh + offs_d * stride_od
-    out_ptrs = Out + off_o
-    tl.store(out_ptrs, acc)
-    return
-
-
-@torch.no_grad()
-def token_att_fwd2_int8v(prob, v, v_scale, out, Req_to_tokens, B_req_idx, B_Start_Loc, B_Seqlen, max_len_in_batch):
-    if max_len_in_batch < 512:
-        BLOCK = triton.next_power_of_2(max_len_in_batch)
-    else:
-        BLOCK = 512
-    batch, head = B_req_idx.shape[0], prob.shape[0]
-    grid = (batch, head)
-    num_warps = 4
-    dim = v.shape[-1]
-    kv_group_num = prob.shape[0] // v.shape[1]
-
-    _fwd_kernel_token_att2_int8v[grid](
-        prob,
-        v,
-        v_scale,
-        out,
-        Req_to_tokens,
-        B_req_idx,
-        B_Start_Loc,
-        B_Seqlen,
-        Req_to_tokens.stride(0),
-        Req_to_tokens.stride(1),
-        prob.stride(0),
-        prob.stride(1),
-        v.stride(0),
-        v.stride(1),
-        v.stride(2),
-        v_scale.stride(0),
-        v_scale.stride(1),
-        v_scale.stride(2),
-        out.stride(0),
-        out.stride(1),
-        out.stride(2),
-        kv_group_num=kv_group_num,
-        BLOCK_DMODEL=dim,
-        BLOCK_N=BLOCK,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
-
-
-def torch_att(V, P, bs, seqlen, num_head, head_dim):
-    V = V.view(bs, seqlen, num_head, head_dim).transpose(1, 2)
-    P = P.reshape(num_head, bs, 1, seqlen).transpose(0, 1)
-    out = torch.matmul(P, V)
-
-    return out
diff --git a/lightllm/models/llava/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/llava/layer_weights/pre_and_post_layer_weight.py
index b4c070a1e6..3afcfb0a71 100644
--- a/lightllm/models/llava/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/llava/layer_weights/pre_and_post_layer_weight.py
@@ -1,5 +1,3 @@
-import torch
-import numpy as np
 from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
 
 
@@ -14,8 +12,8 @@ def rename_weight_keys(weights):
 
 
 class LlavaPreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
         return
 
     def load_hf_weights(self, weights):
diff --git a/lightllm/models/minicpm/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/minicpm/layer_weights/pre_and_post_layer_weight.py
index 0952468d0f..45023bdf8f 100644
--- a/lightllm/models/minicpm/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/minicpm/layer_weights/pre_and_post_layer_weight.py
@@ -3,8 +3,8 @@
 
 
 class MiniCPMPreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
         hidden_size = self.network_config_["hidden_size"]
         dim_model_base = self.network_config_.get("dim_model_base", hidden_size)
         self.lm_head_scale = hidden_size / dim_model_base
diff --git a/lightllm/models/minicpm/layer_weights/transformer_layer_weight.py b/lightllm/models/minicpm/layer_weights/transformer_layer_weight.py
index 2bc5078382..c37b524fde 100755
--- a/lightllm/models/minicpm/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/minicpm/layer_weights/transformer_layer_weight.py
@@ -3,8 +3,8 @@
 
 
 class MiniCPMTransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _parse_config(self):
diff --git a/lightllm/models/mistral/layer_infer/transformer_layer_infer.py b/lightllm/models/mistral/layer_infer/transformer_layer_infer.py
index 59eef6daa7..d115c30ec1 100755
--- a/lightllm/models/mistral/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/mistral/layer_infer/transformer_layer_infer.py
@@ -4,7 +4,7 @@
 class MistralTransformerLayerInfer(LlamaTransformerLayerInfer):
     """ """
 
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         self.head_dim_ = network_config.get("head_dim", self.head_dim_)
         return
diff --git a/lightllm/models/mistral/model.py b/lightllm/models/mistral/model.py
index d32f51ae78..f09525c59f 100644
--- a/lightllm/models/mistral/model.py
+++ b/lightllm/models/mistral/model.py
@@ -1,5 +1,3 @@
-import os
-import json
 import torch
 from lightllm.models.registry import ModelRegistry
 from lightllm.common.basemodel import TpPartBaseModel
@@ -8,7 +6,6 @@
 from lightllm.models.llama.layer_infer.pre_layer_infer import LlamaPreLayerInfer
 from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.models.llama.flashattention_infer_struct import FlashAttentionStateInfo
 from lightllm.models.mistral.layer_infer.transformer_layer_infer import MistralTransformerLayerInfer
 from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class
 from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num
@@ -43,10 +40,6 @@ def _init_custom(self):
         self._init_to_get_rotary()
         return
 
-    def _init_inferstate_cls(self):
-        if get_env_start_args().enable_fa3:
-            self.infer_state_class = FlashAttentionStateInfo
-
     def _init_mem_manager(self):
         # Dealing with head_dim_!=n_embed // num_attention_heads scenarios, such as mistral 13B
         head_dim = self.config["hidden_size"] // self.config["num_attention_heads"]
diff --git a/lightllm/models/mistral/triton_kernel/context_flashattention_nopad.py b/lightllm/models/mistral/triton_kernel/context_flashattention_nopad.py
deleted file mode 100644
index abcaf02b51..0000000000
--- a/lightllm/models/mistral/triton_kernel/context_flashattention_nopad.py
+++ /dev/null
@@ -1,228 +0,0 @@
-import torch
-
-import triton
-import triton.language as tl
-import math
-import torch.nn.functional as F
-
-
-@triton.jit
-def _fwd_kernel(
-    Q,
-    K,
-    V,
-    sm_scale,
-    B_Start_Loc,
-    B_Seqlen,  # B_LOC 内部记录每个batch 输入的真实位置， B_SEQ_len 记录当前输入的真实长度
-    Out,
-    stride_qbs,
-    stride_qh,
-    stride_qd,
-    stride_kbs,
-    stride_kh,
-    stride_kd,
-    stride_vbs,
-    stride_vh,
-    stride_vd,
-    stride_obs,
-    stride_oh,
-    stride_od,
-    kv_group_num,
-    sliding_window,
-    BLOCK_M: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-):
-    cur_batch = tl.program_id(0)
-    cur_head = tl.program_id(1)
-    start_m = tl.program_id(2)
-
-    cur_kv_head = cur_head // kv_group_num
-
-    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
-    cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
-
-    block_start_loc = BLOCK_M * start_m
-
-    # initialize offsets
-    offs_n = tl.arange(0, BLOCK_N)
-    offs_d = tl.arange(0, BLOCK_DMODEL)
-    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    off_q = (
-        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs
-        + cur_head * stride_qh
-        + offs_d[None, :] * stride_qd
-    )
-    off_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd
-    off_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] * stride_vd
-
-    q = tl.load(Q + off_q, mask=offs_m[:, None] < cur_batch_seq_len, other=0.0)
-
-    k_ptrs = K + off_k
-    v_ptrs = V + off_v
-
-    # initialize pointer to m and l
-    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
-    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
-    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
-
-    block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0)
-
-    for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
-        start_n = tl.multiple_of(start_n, BLOCK_N)
-        # -- compute qk ----
-        k = tl.load(
-            k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
-            mask=(start_n + offs_n[None, :]) < cur_batch_seq_len,
-            other=0.0,
-        )
-        # mask = tl.load(mask_ptrs + start_n, mask=start_n + offs_n < cur_batch_end_loc, other=0.0)
-
-        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
-        qk += tl.dot(q, k)
-        qk *= sm_scale
-        # [SYM] mask outside of windows
-        qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf"))
-        qk = tl.where((start_n + offs_n[None, :]) > (offs_m[:, None] - sliding_window), qk, float("-inf"))
-
-        # -- compute m_ij, p, l_ij
-        m_ij = tl.max(qk, 1)
-        p = tl.exp(qk - m_ij[:, None])
-        l_ij = tl.sum(p, 1)
-        # -- update m_i and l_i
-        m_i_new = tl.maximum(m_i, m_ij)
-        alpha = tl.exp(m_i - m_i_new)
-        beta = tl.exp(m_ij - m_i_new)
-        l_i_new = alpha * l_i + beta * l_ij
-        # -- update output accumulator --
-        # scale p
-        p_scale = beta / l_i_new
-        p = p * p_scale[:, None]
-        # scale acc
-        acc_scale = l_i / l_i_new * alpha
-        acc = acc * acc_scale[:, None]
-        # update acc
-        v = tl.load(
-            v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
-            mask=(start_n + offs_n[:, None]) < cur_batch_seq_len,
-            other=0.0,
-        )
-
-        p = p.to(v.dtype)
-        acc += tl.dot(p, v)
-        # update m_i and l_i
-        l_i = l_i_new
-        m_i = m_i_new
-    # initialize pointers to output
-    off_o = (
-        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
-        + cur_head * stride_oh
-        + offs_d[None, :] * stride_od
-    )
-    out_ptrs = Out + off_o
-    tl.store(out_ptrs, acc, mask=offs_m[:, None] < cur_batch_seq_len)
-    return
-
-
-@torch.no_grad()
-def context_attention_fwd(q, k, v, o, b_start_loc, b_seq_len, max_input_len, sliding_window):
-    BLOCK = 128
-    # shape constraints
-    Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
-    assert Lq == Lk and Lk == Lv
-    assert Lk in {16, 32, 64, 128}
-
-    sm_scale = 1.0 / (Lq ** 0.5)  # 计算scale系数
-    batch, head = b_seq_len.shape[0], q.shape[1]
-    kv_group_num = q.shape[1] // k.shape[1]
-
-    grid = (batch, head, triton.cdiv(max_input_len, BLOCK))  # batch, head,
-
-    num_warps = 4 if Lk <= 64 else 8
-    _fwd_kernel[grid](
-        q,
-        k,
-        v,
-        sm_scale,
-        b_start_loc,
-        b_seq_len,
-        o,
-        q.stride(0),
-        q.stride(1),
-        q.stride(2),
-        k.stride(0),
-        k.stride(1),
-        k.stride(2),
-        v.stride(0),
-        v.stride(1),
-        v.stride(2),
-        o.stride(0),
-        o.stride(1),
-        o.stride(2),
-        kv_group_num=kv_group_num,
-        sliding_window=sliding_window,
-        BLOCK_M=BLOCK,
-        BLOCK_DMODEL=Lk,
-        BLOCK_N=BLOCK,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
-
-
-def torch_att(xq, xk, xv, bs, seqlen, num_head, head_dim):
-    xq = xq.view(bs, seqlen, num_head, head_dim)
-    xk = xk.view(bs, seqlen, num_head, head_dim)
-    xv = xv.view(bs, seqlen, num_head, head_dim)
-    mask = torch.tril(torch.ones(seqlen, seqlen), diagonal=0).unsqueeze(0).unsqueeze(0).cuda()
-    mask[mask == 0.0] = -100000000.0
-    mask = mask.repeat(bs, num_head, 1, 1)
-    keys = xk
-    values = xv
-    xq = xq.transpose(1, 2)
-    keys = keys.transpose(1, 2)
-    values = values.transpose(1, 2)
-    scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(head_dim)
-    scores = F.softmax(scores.float() + mask, dim=-1).type_as(xq)
-    output = torch.matmul(scores, values).transpose(1, 2).contiguous().reshape(-1, num_head, head_dim)
-    return output
-
-
-def test():
-    import torch
-
-    Z, H, N_CTX, D_HEAD = 4, 6, 1024, 128
-    dtype = torch.float16
-    Z = 3
-    q = torch.empty((Z * N_CTX, H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.1, std=0.2)
-    k = torch.empty((Z * N_CTX, H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.4, std=0.2)
-    v = torch.empty((Z * N_CTX, H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.3, std=0.2)
-    o = torch.empty((Z * N_CTX, H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.3, std=0.2)
-
-    max_input_len = N_CTX
-    Z = 4
-    b_start_loc = torch.zeros((Z,), dtype=torch.int32, device="cuda")
-    b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda")
-
-    b_seq_len[0] = 512
-    b_seq_len[1] = 1024
-    b_seq_len[2] = 512
-    b_seq_len[3] = 1024
-
-    for i in range(1, Z):
-        b_start_loc[i] = b_start_loc[i - 1] + b_seq_len[i - 1]
-
-    torch_out = []
-    start = 0
-    for i in range(Z):
-        end = start + b_seq_len[i]
-        torch_o = torch_att(q[start:end], k[start:end], v[start:end], 1, b_seq_len[i], H, D_HEAD)
-        start = end
-        torch_out.append(torch_o)
-    torch_out = torch.cat(torch_out, dim=0)
-    context_attention_fwd(q, k, v, o, b_start_loc, b_seq_len, max_input_len, 10)
-    print(o.shape, torch_out.shape)
-
-    print("max ", torch.max(torch.abs(torch_out - o)))
-    print("mean ", torch.mean(torch.abs(torch_out - o)))
-    assert torch.allclose(torch_out, o, atol=1e-2, rtol=0)
diff --git a/lightllm/models/mistral/triton_kernel/init_att_sliding_window_info.py b/lightllm/models/mistral/triton_kernel/init_att_sliding_window_info.py
deleted file mode 100644
index a60fe970b3..0000000000
--- a/lightllm/models/mistral/triton_kernel/init_att_sliding_window_info.py
+++ /dev/null
@@ -1,45 +0,0 @@
-import torch
-
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def _fwd_kernel_init_att_window_info(
-    b_seq_len,
-    b_att_seq_len,
-    batch_size,
-    sliding_window,
-    BLOCK_SIZE: tl.constexpr,
-):
-    cur_index = tl.program_id(0)
-    cur_start = cur_index * BLOCK_SIZE
-    offsets = cur_start + tl.arange(0, BLOCK_SIZE)
-    mask = offsets < batch_size
-
-    cur_seq_len = tl.load(b_seq_len + offsets, mask=mask)
-    b_att_seq_len_data = tl.minimum(cur_seq_len, sliding_window)
-
-    tl.store(b_att_seq_len + offsets, b_att_seq_len_data, mask=mask)
-    return
-
-
-@torch.no_grad()
-def init_att_window_info_fwd(batch_size, b_seq_len, b_att_seq_len, sliding_window):
-    # shape constraints
-    assert batch_size == b_seq_len.shape[0] == b_att_seq_len.shape[0]
-
-    BLOCK_SIZE = 32
-    num_warps = 1
-    grid = (triton.cdiv(batch_size, BLOCK_SIZE),)
-
-    _fwd_kernel_init_att_window_info[grid](
-        b_seq_len,
-        b_att_seq_len,
-        batch_size=batch_size,
-        sliding_window=sliding_window,
-        BLOCK_SIZE=BLOCK_SIZE,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
diff --git a/lightllm/models/mistral/triton_kernel/token_attention_softmax_and_reducev.py b/lightllm/models/mistral/triton_kernel/token_attention_softmax_and_reducev.py
deleted file mode 100644
index bf9928f987..0000000000
--- a/lightllm/models/mistral/triton_kernel/token_attention_softmax_and_reducev.py
+++ /dev/null
@@ -1,132 +0,0 @@
-import torch
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def _fwd_kernel(
-    Logics,
-    V,
-    Out,
-    Req_to_tokens,
-    B_req_idx,
-    B_Start_Loc,
-    B_Seqlen,
-    B_Att_Start_Loc,
-    B_Att_Seqlen,
-    stride_logic_h,
-    stride_logic_bs,
-    stride_vbs,
-    stride_vh,
-    stride_vd,
-    stride_obs,
-    stride_oh,
-    stride_od,
-    stride_req_to_token_b,
-    stride_req_to_token_s,
-    other_kv_index,  # 避免读取到nan的数据
-    kv_group_num,
-    sliding_window,
-    BLOCK_DMODEL: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-):
-    cur_batch = tl.program_id(0)
-    cur_head = tl.program_id(1)
-
-    cur_kv_head = cur_head // kv_group_num
-
-    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
-    cur_batch_start_loc = tl.load(B_Att_Start_Loc + cur_batch)  # new index
-    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
-    cur_att_seq_len = tl.load(B_Att_Seqlen + cur_batch)  # new index
-    cur_cache_start_loc = tl.maximum(cur_batch_seq_len - sliding_window, 0)  # new index
-
-    offs_n = tl.arange(0, BLOCK_N)  # [64]
-    offs_d = tl.arange(0, BLOCK_DMODEL)  # [D]
-
-    off_v = cur_kv_head * stride_vh + offs_d[None, :] * stride_vd  # [1, D]
-    v_ptrs = V + off_v
-
-    e_max = float("-inf")
-    e_sum = 0.0
-    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)  # [D]
-
-    for start_n in range(0, cur_att_seq_len, BLOCK_N):
-        start_n = tl.multiple_of(start_n, BLOCK_N)  # check
-        v_index = tl.load(
-            Req_to_tokens
-            + cur_batch_req_idx * stride_req_to_token_b
-            + (cur_cache_start_loc + start_n + offs_n) * stride_req_to_token_s,
-            mask=(start_n + offs_n) < cur_att_seq_len,
-            other=other_kv_index,
-        )  # [64]
-
-        qk = tl.load(
-            Logics + cur_head * stride_logic_h + (cur_batch_start_loc + start_n + offs_n) * stride_logic_bs,
-            mask=(start_n + offs_n) < cur_att_seq_len,
-            other=float("-inf"),
-        )  # [64]
-
-        n_e_max = tl.maximum(tl.max(qk, 0), e_max)
-        old_scale = tl.exp(e_max - n_e_max)
-        p = tl.exp(qk - n_e_max)
-        e_sum = e_sum * old_scale + tl.sum(p, 0)
-        v = tl.load(v_ptrs + v_index[:, None] * stride_vbs)  # [1, D] + [64, 1] = [64, D]
-        acc = acc * old_scale + tl.sum(p[:, None] * v, 0)  # [64, 1] * [64, D] = [64, D] -> [D]
-        e_max = n_e_max
-
-    acc = acc / e_sum
-    off_o = cur_batch * stride_obs + cur_head * stride_oh + offs_d * stride_od
-    out_ptrs = Out + off_o
-    tl.store(out_ptrs, acc)
-    return
-
-
-@torch.no_grad()
-def token_softmax_reducev_fwd(
-    logics,
-    v,
-    o,
-    req_to_tokens,
-    b_req_idx,
-    b_start_loc,
-    b_seq_len,
-    b_att_start_loc,
-    b_att_seq_len,
-    sliding_window,
-):
-    BLOCK = 64
-    batch, head = b_seq_len.shape[0], logics.shape[0]
-    grid = (batch, head)
-    kv_group_num = logics.shape[0] // v.shape[1]
-
-    num_warps = 1
-    _fwd_kernel[grid](
-        logics,
-        v,
-        o,
-        req_to_tokens,
-        b_req_idx,
-        b_start_loc,
-        b_seq_len,
-        b_att_start_loc,
-        b_att_seq_len,
-        logics.stride(0),
-        logics.stride(1),
-        v.stride(0),
-        v.stride(1),
-        v.stride(2),
-        o.stride(0),
-        o.stride(1),
-        o.stride(2),
-        req_to_tokens.stride(0),
-        req_to_tokens.stride(1),
-        0,
-        kv_group_num,
-        sliding_window,
-        BLOCK_DMODEL=v.shape[-1],
-        BLOCK_N=BLOCK,
-        num_warps=num_warps,
-        num_stages=3,
-    )
-    return
diff --git a/lightllm/models/mistral_mtp/layer_infer/post_layer_infer.py b/lightllm/models/mistral_mtp/layer_infer/post_layer_infer.py
index 5eac249bad..f890fbf663 100644
--- a/lightllm/models/mistral_mtp/layer_infer/post_layer_infer.py
+++ b/lightllm/models/mistral_mtp/layer_infer/post_layer_infer.py
@@ -4,6 +4,6 @@
 class MistralMTPPostLayerInfer(LlamaPostLayerInfer):
     """ """
 
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         return
diff --git a/lightllm/models/mistral_mtp/layer_infer/pre_layer_infer.py b/lightllm/models/mistral_mtp/layer_infer/pre_layer_infer.py
index 25bea1aa60..96a15d18a6 100644
--- a/lightllm/models/mistral_mtp/layer_infer/pre_layer_infer.py
+++ b/lightllm/models/mistral_mtp/layer_infer/pre_layer_infer.py
@@ -7,8 +7,8 @@
 class MistralMTPPreLayerInfer(LlamaPreLayerInfer):
     """ """
 
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         return
 
     def _mtp_context_forward(
@@ -19,22 +19,10 @@ def _mtp_context_forward(
             input_embdings.shape[0] == tgt_embdings.shape[0]
         ), f"shape {input_embdings.shape} != shape {tgt_embdings.shape}"
 
-        layer_weight.enorm_weight_.rmsnorm_forward(
-            input=input_embdings,
-            eps=self.eps_,
-            out=input_embdings,
-        )
+        layer_weight.enorm_weight_(input=input_embdings, eps=self.eps_, out=input_embdings)
 
-        tgt_embdings = layer_weight.final_norm_weight_.rmsnorm_forward(
-            input=tgt_embdings,
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
-        )
-        layer_weight.hnorm_weight_.rmsnorm_forward(
-            input=tgt_embdings,
-            eps=self.eps_,
-            out=tgt_embdings,
-        )
+        tgt_embdings = layer_weight.final_norm_weight_(input=tgt_embdings, eps=self.eps_, alloc_func=self.alloc_tensor)
+        layer_weight.hnorm_weight_(input=tgt_embdings, eps=self.eps_, out=tgt_embdings)
 
         cat_embdings = torch.cat((input_embdings, tgt_embdings), dim=-1)
 
@@ -47,22 +35,10 @@ def _mtp_token_forward(
         tgt_embdings = infer_state.mtp_draft_input_hiddens
         assert input_embdings.shape[0] == tgt_embdings.shape[0]
 
-        layer_weight.enorm_weight_.rmsnorm_forward(
-            input=input_embdings,
-            eps=self.eps_,
-            out=input_embdings,
-        )
+        layer_weight.enorm_weight_(input=input_embdings, eps=self.eps_, out=input_embdings)
 
-        tgt_embdings = layer_weight.final_norm_weight_.rmsnorm_forward(
-            input=tgt_embdings,
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
-        )
-        layer_weight.hnorm_weight_.rmsnorm_forward(
-            input=tgt_embdings,
-            eps=self.eps_,
-            out=tgt_embdings,
-        )
+        tgt_embdings = layer_weight.final_norm_weight_(input=tgt_embdings, eps=self.eps_, alloc_func=self.alloc_tensor)
+        layer_weight.hnorm_weight_(input=tgt_embdings, eps=self.eps_, out=tgt_embdings)
 
         cat_embdings = torch.cat((input_embdings, tgt_embdings), dim=-1)
 
diff --git a/lightllm/models/mistral_mtp/layer_infer/transformer_layer_infer.py b/lightllm/models/mistral_mtp/layer_infer/transformer_layer_infer.py
index 5724f32af9..6d72ae2c38 100644
--- a/lightllm/models/mistral_mtp/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/mistral_mtp/layer_infer/transformer_layer_infer.py
@@ -10,8 +10,8 @@
 
 
 class MistralMTPTransformerLayerInfer(MistralTransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         return
 
     def context_forward(self, input_embdings, infer_state: InferStateInfo, layer_weight):
diff --git a/lightllm/models/mistral_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/mistral_mtp/layer_weights/pre_and_post_layer_weight.py
index 2fbc89cfd0..a65250b16d 100644
--- a/lightllm/models/mistral_mtp/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/mistral_mtp/layer_weights/pre_and_post_layer_weight.py
@@ -2,15 +2,15 @@
 from lightllm.common.basemodel.layer_weights.meta_weights import (
     EmbeddingWeight,
     LMHeadWeight,
-    NoTpNormWeight,
+    RMSNormWeight,
     ROWMMWeight,
 )
 
 
 class MistralMTPPreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
-
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
+        hidden_size = network_config["hidden_size"]
         self.eh_proj_weight_ = ROWMMWeight(
             weight_names="mtp.eh_proj.weight",
             data_type=self.data_type_,
@@ -19,12 +19,13 @@ def __init__(self, data_type, network_config, mode):
             tp_rank=0,
             tp_world_size=1,
         )
-        self.enorm_weight_ = NoTpNormWeight(
+        self.enorm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="mtp.enorm.weight",
             data_type=self.data_type_,
-            bias_name=None,
         )
-        self.hnorm_weight_ = NoTpNormWeight(
+        self.hnorm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="mtp.hnorm.weight",
             data_type=self.data_type_,
             bias_name=None,
@@ -32,5 +33,5 @@ def __init__(self, data_type, network_config, mode):
 
         self.wte_weight_: EmbeddingWeight = None
         self.lm_head_weight_: LMHeadWeight = None
-        self.final_norm_weight_: NoTpNormWeight = None
+        self.final_norm_weight_: RMSNormWeight = None
         return
diff --git a/lightllm/models/mistral_mtp/layer_weights/transformer_layer_weight.py b/lightllm/models/mistral_mtp/layer_weights/transformer_layer_weight.py
index 6607dbb704..2cbc6cf585 100644
--- a/lightllm/models/mistral_mtp/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/mistral_mtp/layer_weights/transformer_layer_weight.py
@@ -1,12 +1,16 @@
 from lightllm.common.basemodel import TransformerLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import ROWMMWeight, COLMMWeight, NoTpNormWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import ROWMMWeight, COLMMWeight, RMSNormWeight
 
 
 class MistralMTPTransformerLayerWeight(TransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
+    def _parse_config(self):
+        self.n_embed = self.network_config_["hidden_size"]
+        self.n_inter = self.network_config_["intermediate_size"]
+
     def _init_weight_names(self):
         self._gate_weight_name = f"mtp.layers.{self.layer_num_}.mlp.gate_proj.weight"
         self._gate_bias_name = None
@@ -24,23 +28,27 @@ def _init_weight(self):
 
     def _init_ffn(self):
         self.gate_up_proj = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_inter, self.n_inter],
             weight_names=[self._gate_weight_name, self._up_weight_name],
             data_type=self.data_type_,
             bias_names=[self._gate_bias_name, self._up_bias_name],
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="gate_up_proj",
+            quant_method=self.get_quant_method("gate_up_proj"),
         )
         self.down_proj = COLMMWeight(
+            in_dim=self.n_inter,
+            out_dims=[self.n_embed],
             weight_names=self._down_weight_name,
             data_type=self.data_type_,
             bias_names=self._down_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="down_proj",
+            quant_method=self.get_quant_method("down_proj"),
         )
 
     def _init_norm(self):
-        self.ffn_norm_weight_ = NoTpNormWeight(
-            self._ffn_norm_weight_name, self.data_type_, bias_name=self._ffn_norm_bias_name
+        hidden_size = self.network_config_["hidden_size"]
+        self.ffn_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
+            weight_name=self._ffn_norm_weight_name,
+            data_type=self.data_type_,
+            bias_name=self._ffn_norm_bias_name,
         )
diff --git a/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py b/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py
index a60375688b..44e66cff2d 100644
--- a/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py
@@ -3,14 +3,13 @@
 import torch.nn.functional as F
 from lightllm.common.basemodel.infer_struct import InferStateInfo
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
-from lightllm.models.mistral.layer_infer.transformer_layer_infer import MistralTransformerLayerInfer
 from lightllm.models.mixtral.layer_infer._custom_ops import fused_topk
 from lightllm.models.mixtral.layer_weights.transformer_layer_weight import MixtralTransformerLayerWeight
 
 
 class MixtralTransformerLayerInfer(LlamaTransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         self.num_local_experts = network_config["num_local_experts"]
         self.num_experts_per_tok = network_config["num_experts_per_tok"]
         self.renormalize = True
diff --git a/lightllm/models/mixtral/layer_weights/transformer_layer_weight.py b/lightllm/models/mixtral/layer_weights/transformer_layer_weight.py
index f425ad08ba..51c62fd4cb 100644
--- a/lightllm/models/mixtral/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/mixtral/layer_weights/transformer_layer_weight.py
@@ -2,18 +2,18 @@
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.envs_utils import enable_env_vars
 from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import ROWMMWeight, FusedMoeWeightEP, create_tp_moe_wegiht_obj
+from lightllm.common.basemodel.layer_weights.meta_weights import ROWMMWeight, FusedMoeWeight
+from lightllm.utils.envs_utils import get_env_start_args
 
 logger = init_logger(__name__)
 
 
 class MixtralTransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
         super().__init__(
             layer_num,
             data_type,
             network_config,
-            mode,
             quant_cfg=quant_cfg,
         )
         return
@@ -32,36 +32,29 @@ def _init_ffn(self):
 
     def _init_moe(self):
         inter_size = self.network_config_["intermediate_size"]
-        split_inter_size = inter_size // self.tp_world_size_
 
         self.moe_gate = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_routed_experts],
             weight_names=self.moe_gate_weight_name,
             data_type=self.data_type_,
             bias_names=self.moe_gate_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="moe_gate",
+            quant_method=self.get_quant_method("moe_gate"),
             tp_rank=0,
             tp_world_size=1,  # no tensor parallelism
         )
-
-        moe_mode = os.getenv("MOE_MODE", "TP")
-        assert moe_mode in ["TP"], f"Unsupported moe mode: {moe_mode}"
-
-        if moe_mode == "TP":
-            self.experts = create_tp_moe_wegiht_obj(
-                gate_proj_name="w1",
-                down_proj_name="w2",
-                up_proj_name="w3",
-                e_score_correction_bias_name="",
-                weight_prefix=f"model.layers.{self.layer_num_}.block_sparse_moe.experts",
-                n_routed_experts=self.n_routed_experts,
-                split_inter_size=split_inter_size,
-                data_type=self.data_type_,
-                network_config=self.network_config_,
-                layer_num=self.layer_num_,
-                quant_cfg=self.quant_cfg,
-                num_fused_shared_experts=0,
-            )
-        else:
-            raise ValueError(f"Unsupported moe mode: {moe_mode}")
+        assert get_env_start_args().enable_ep_moe, "Mixtral only support tp mode."
+        self.experts = FusedMoeWeight(
+            gate_proj_name="w1",
+            down_proj_name="w2",
+            up_proj_name="w3",
+            e_score_correction_bias_name="",
+            weight_prefix=f"model.layers.{self.layer_num_}.block_sparse_moe.experts",
+            n_routed_experts=self.n_routed_experts,
+            hidden_size=self.n_embed,
+            moe_intermediate_size=inter_size,
+            data_type=self.data_type_,
+            quant_method=self.quant_cfg.get_quant_method(self.layer_num_, "fused_moe"),
+            layer_num=self.layer_num_,
+            network_config=self.network_config_,
+        )
diff --git a/lightllm/models/phi3/layer_infer/transformer_layer_infer.py b/lightllm/models/phi3/layer_infer/transformer_layer_infer.py
index 806c59365b..fd3d05e426 100755
--- a/lightllm/models/phi3/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/phi3/layer_infer/transformer_layer_infer.py
@@ -1,11 +1,5 @@
-import torch
-from functools import partial
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
 from lightllm.models.phi3.triton_kernel.rotary_emb import rotary_emb_fwd
-from lightllm.models.phi3.triton_kernel.context_flashattention_nopad import (
-    context_attention_fwd,
-)
-from lightllm.models.phi3.triton_kernel.destindex_copy_kv import destindex_copy_kv
 from lightllm.models.phi3.layer_weights.transformer_layer_weight import Phi3TransformerLayerWeight
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
 
@@ -13,14 +7,8 @@
 class Phi3TransformerLayerInfer(LlamaTransformerLayerInfer):
     """ """
 
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
-        return
-
-    def _bind_attention(self):
-        self._context_attention_kernel = partial(Phi3TransformerLayerInfer._context_attention_kernel, self)
-        self._copy_kv_to_mem_cache = partial(Phi3TransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-        self._token_attention_kernel = partial(Phi3TransformerLayerInfer._token_decode_attention_flashdecoding, self)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         return
 
     def _get_qkv(self, input_emb, infer_state: LlamaInferStateInfo, layer_weight: Phi3TransformerLayerWeight):
@@ -35,44 +23,3 @@ def _get_qkv(self, input_emb, infer_state: LlamaInferStateInfo, layer_weight: Ph
             infer_state.position_sin,
         )
         return q, cache_kv
-
-    def _copy_kv_to_mem_cache_normal(self, buffer, mem_index, mem_manager):
-        destindex_copy_kv(buffer, mem_index, mem_manager.kv_buffer[self.layer_num_])
-        return
-
-    def _context_attention_kernel(
-        self, q, kv, infer_state: LlamaInferStateInfo, layer_weight, out=None
-    ) -> torch.Tensor:
-        o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        context_attention_fwd(
-            q.view(-1, self.tp_q_head_num_, self.head_dim_),
-            kv[:, 0 : self.tp_k_head_num_, :],
-            kv[:, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :],
-            o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_),
-            infer_state.b_req_idx,
-            infer_state.b_start_loc,
-            infer_state.b_seq_len,
-            infer_state.b_ready_cache_len,
-            infer_state.max_len_in_batch,
-            infer_state.req_manager.req_to_token_indexs,
-        )
-        return o_tensor
-
-    def _token_decode_attention_flashdecoding(self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None):
-        from lightllm.models.phi3.triton_kernel.flash_decoding import token_decode_attention_flash_decoding
-
-        cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :]
-        cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
-            :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ]
-        return token_decode_attention_flash_decoding(
-            q,
-            infer_state,
-            self.tp_q_head_num_,
-            self.head_dim_,
-            cache_k,
-            cache_v,
-            out=out,
-            alloc_tensor_func=self.alloc_tensor,
-        )
diff --git a/lightllm/models/phi3/layer_weights/transformer_layer_weight.py b/lightllm/models/phi3/layer_weights/transformer_layer_weight.py
index 91b2730917..db4906c19a 100755
--- a/lightllm/models/phi3/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/phi3/layer_weights/transformer_layer_weight.py
@@ -6,8 +6,8 @@
 
 
 class Phi3TransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def load_hf_weights(self, weights):
diff --git a/lightllm/models/phi3/triton_kernel/context_flashattention_nopad.py b/lightllm/models/phi3/triton_kernel/context_flashattention_nopad.py
deleted file mode 100644
index ee04c3367b..0000000000
--- a/lightllm/models/phi3/triton_kernel/context_flashattention_nopad.py
+++ /dev/null
@@ -1,433 +0,0 @@
-import torch
-
-import triton
-import triton.language as tl
-import math
-import torch.nn.functional as F
-
-from lightllm.utils.device_utils import is_tesla
-
-
-@triton.jit
-def _fwd_kernel(
-    Q,
-    K,
-    V,
-    sm_scale,
-    B_Start_Loc,
-    B_Seqlen,  # B_LOC 内部记录每个batch 输入的真实位置， B_SEQ_len 记录当前输入的真实长度
-    Out,
-    Req_to_tokens,
-    B_req_idx,
-    stride_qbs,
-    stride_qh,
-    stride_qd,
-    stride_kbs,
-    stride_kh,
-    stride_kd,
-    stride_vbs,
-    stride_vh,
-    stride_vd,
-    stride_obs,
-    stride_oh,
-    stride_od,
-    stride_req_to_tokens_b,
-    stride_req_to_tokens_s,
-    kv_group_num,
-    b_prompt_cache_len,
-    head_dim: tl.constexpr,
-    BLOCK_M: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-):
-    cur_batch = tl.program_id(0)
-    cur_head = tl.program_id(1)
-    start_m = tl.program_id(2)
-
-    cur_kv_head = cur_head // kv_group_num
-
-    cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
-    prompt_cache_len = tl.load(b_prompt_cache_len + cur_batch)
-    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) - prompt_cache_len
-    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
-
-    block_start_loc = BLOCK_M * start_m
-
-    # initialize offsets
-    offs_n = tl.arange(0, BLOCK_N)
-    offs_d = tl.arange(0, BLOCK_DMODEL)
-    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    off_q = (
-        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs
-        + cur_head * stride_qh
-        + offs_d[None, :] * stride_qd
-    )
-
-    q = tl.load(Q + off_q, mask=(offs_m[:, None] < cur_batch_seq_len) & (offs_d[None, :] < head_dim), other=0.0)
-
-    # initialize pointer to m and l
-    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
-    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
-    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
-
-    block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0)
-    block_end_loc = tl.minimum((start_m + 1) * BLOCK_M + prompt_cache_len, cur_batch_seq_len + prompt_cache_len)
-
-    for start_n in range(0, block_mask * block_end_loc, BLOCK_N):
-        start_n = tl.multiple_of(start_n, BLOCK_N)
-        # -- compute qk ----
-        kv_loc = tl.load(
-            Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * (start_n + offs_n),
-            mask=(start_n + offs_n) < block_end_loc,
-            other=0,
-        ).to(tl.int64)
-        off_k = kv_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd
-        k = tl.load(
-            K + off_k, mask=((start_n + offs_n[None, :]) < block_end_loc) & (offs_d[:, None] < head_dim), other=0.0
-        )
-
-        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
-        qk += tl.dot(q, k)
-        qk *= sm_scale
-        qk = tl.where(offs_m[:, None] + prompt_cache_len >= start_n + offs_n[None, :], qk, float("-100000000.0"))
-
-        # -- compute m_ij, p, l_ij
-        m_ij = tl.max(qk, 1)
-        p = tl.exp(qk - m_ij[:, None])
-        l_ij = tl.sum(p, 1)
-        # -- update m_i and l_i
-        m_i_new = tl.maximum(m_i, m_ij)
-        alpha = tl.exp(m_i - m_i_new)
-        beta = tl.exp(m_ij - m_i_new)
-        l_i_new = alpha * l_i + beta * l_ij
-        # -- update output accumulator --
-        # scale p
-        p_scale = beta / l_i_new
-        p = p * p_scale[:, None]
-        # scale acc
-        acc_scale = l_i / l_i_new * alpha
-        acc_scale = tl.where(offs_m + prompt_cache_len >= start_n, acc_scale, 1.0)
-        acc = acc * acc_scale[:, None]
-        # update acc
-        off_v = kv_loc[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] * stride_vd
-        v = tl.load(
-            V + off_v, mask=((start_n + offs_n[:, None]) < block_end_loc) & (offs_d[None, :] < head_dim), other=0.0
-        )
-        p = p.to(v.dtype)
-        acc += tl.dot(p, v)
-        # update m_i and l_i
-        l_i = l_i_new
-        m_i = m_i_new
-    # initialize pointers to output
-    off_o = (
-        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
-        + cur_head * stride_oh
-        + offs_d[None, :] * stride_od
-    )
-    out_ptrs = Out + off_o
-    tl.store(out_ptrs, acc, mask=(offs_m[:, None] < cur_batch_seq_len) & (offs_d[None, :] < head_dim))
-    return
-
-
-@torch.no_grad()
-def context_attention_fwd(
-    q, k, v, o, b_req_idx, b_start_loc, b_seq_len, b_prompt_cache_len, max_input_len, req_to_token_indexs
-):
-    BLOCK = 128 if not is_tesla() else 64
-    # shape constraints
-    Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
-    assert Lq == Lk and Lk == Lv
-    head_dim = Lq
-    BLOCK_DMODEL = triton.next_power_of_2(head_dim)
-
-    sm_scale = 1.0 / (Lq ** 0.5)  # 计算scale系数
-    batch, head = b_seq_len.shape[0], q.shape[1]
-    kv_group_num = q.shape[1] // k.shape[1]
-
-    grid = (batch, head, triton.cdiv(max_input_len, BLOCK))  # batch, head,
-
-    num_warps = 4 if Lk <= 64 else 8
-    _fwd_kernel[grid](
-        q,
-        k,
-        v,
-        sm_scale,
-        b_start_loc,
-        b_seq_len,
-        o,
-        req_to_token_indexs,
-        b_req_idx,
-        q.stride(0),
-        q.stride(1),
-        q.stride(2),
-        k.stride(0),
-        k.stride(1),
-        k.stride(2),
-        v.stride(0),
-        v.stride(1),
-        v.stride(2),
-        o.stride(0),
-        o.stride(1),
-        o.stride(2),
-        req_to_token_indexs.stride(0),
-        req_to_token_indexs.stride(1),
-        kv_group_num=kv_group_num,
-        b_prompt_cache_len=b_prompt_cache_len,
-        head_dim=head_dim,
-        BLOCK_M=BLOCK,
-        BLOCK_DMODEL=BLOCK_DMODEL,
-        BLOCK_N=BLOCK,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
-
-
-@triton.jit
-def _fwd_kernel_no_prompt_cache(
-    Q,
-    K,
-    V,
-    sm_scale,
-    B_Start_Loc,
-    B_Seqlen,  # B_LOC 内部记录每个batch 输入的真实位置， B_SEQ_len 记录当前输入的真实长度
-    Out,
-    stride_qbs,
-    stride_qh,
-    stride_qd,
-    stride_kbs,
-    stride_kh,
-    stride_kd,
-    stride_vbs,
-    stride_vh,
-    stride_vd,
-    stride_obs,
-    stride_oh,
-    stride_od,
-    kv_group_num,
-    head_dim,
-    BLOCK_M: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-):
-    cur_batch = tl.program_id(0)
-    cur_head = tl.program_id(1)
-    start_m = tl.program_id(2)
-
-    cur_kv_head = cur_head // kv_group_num
-
-    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
-    cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
-
-    block_start_loc = BLOCK_M * start_m
-
-    # initialize offsets
-    offs_n = tl.arange(0, BLOCK_N)
-    offs_d = tl.arange(0, BLOCK_DMODEL)
-    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    off_q = (
-        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs
-        + cur_head * stride_qh
-        + offs_d[None, :] * stride_qd
-    )
-    off_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd
-    off_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] * stride_vd
-
-    q = tl.load(Q + off_q, mask=(offs_m[:, None] < cur_batch_seq_len) & (offs_d[None, :] < head_dim), other=0.0)
-
-    k_ptrs = K + off_k
-    v_ptrs = V + off_v
-
-    # initialize pointer to m and l
-    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
-    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
-    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
-
-    block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0)
-
-    for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
-        start_n = tl.multiple_of(start_n, BLOCK_N)
-        # -- compute qk ----
-        k = tl.load(
-            k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
-            mask=((start_n + offs_n[None, :]) < cur_batch_seq_len) & (offs_d[:, None] < head_dim),
-            other=0.0,
-        )
-        # mask = tl.load(mask_ptrs + start_n, mask=start_n + offs_n < cur_batch_end_loc, other=0.0)
-
-        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
-        qk += tl.dot(q, k)
-        qk *= sm_scale
-        qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf"))
-
-        # -- compute m_ij, p, l_ij
-        m_ij = tl.max(qk, 1)
-        p = tl.exp(qk - m_ij[:, None])
-        l_ij = tl.sum(p, 1)
-        # -- update m_i and l_i
-        m_i_new = tl.maximum(m_i, m_ij)
-        alpha = tl.exp(m_i - m_i_new)
-        beta = tl.exp(m_ij - m_i_new)
-        l_i_new = alpha * l_i + beta * l_ij
-        # -- update output accumulator --
-        # scale p
-        p_scale = beta / l_i_new
-        p = p * p_scale[:, None]
-        # scale acc
-        acc_scale = l_i / l_i_new * alpha
-        acc = acc * acc_scale[:, None]
-        # update acc
-        v = tl.load(
-            v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
-            mask=((start_n + offs_n[:, None]) < cur_batch_seq_len) & (offs_d[None, :] < head_dim),
-            other=0.0,
-        )
-
-        p = p.to(v.dtype)
-        acc += tl.dot(p, v)
-        # update m_i and l_i
-        l_i = l_i_new
-        m_i = m_i_new
-    # initialize pointers to output
-    off_o = (
-        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
-        + cur_head * stride_oh
-        + offs_d[None, :] * stride_od
-    )
-    out_ptrs = Out + off_o
-    tl.store(out_ptrs, acc, mask=(offs_m[:, None] < cur_batch_seq_len) & (offs_d[None, :] < head_dim))
-    return
-
-
-@torch.no_grad()
-def context_attention_fwd_no_prompt_cache(q, k, v, o, b_start_loc, b_seq_len, max_input_len):
-    BLOCK = 128 if not is_tesla() else 64
-    # shape constraints
-    Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
-    assert Lq == Lk and Lk == Lv
-    head_dim = Lq
-    BLOCK_DMODEL = triton.next_power_of_2(head_dim)
-    sm_scale = 1.0 / (Lq ** 0.5)  # 计算scale系数
-    batch, head = b_seq_len.shape[0], q.shape[1]
-    kv_group_num = q.shape[1] // k.shape[1]
-
-    grid = (batch, head, triton.cdiv(max_input_len, BLOCK))  # batch, head,
-
-    num_warps = 4 if Lk <= 64 else 8
-    _fwd_kernel_no_prompt_cache[grid](
-        q,
-        k,
-        v,
-        sm_scale,
-        b_start_loc,
-        b_seq_len,
-        o,
-        q.stride(0),
-        q.stride(1),
-        q.stride(2),
-        k.stride(0),
-        k.stride(1),
-        k.stride(2),
-        v.stride(0),
-        v.stride(1),
-        v.stride(2),
-        o.stride(0),
-        o.stride(1),
-        o.stride(2),
-        kv_group_num=kv_group_num,
-        head_dim=head_dim,
-        BLOCK_M=BLOCK,
-        BLOCK_DMODEL=BLOCK_DMODEL,
-        BLOCK_N=BLOCK,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
-
-
-def torch_att(xq, xk, xv, bs, seqlen, num_head, head_dim, prompt_cache_len):
-    xq = xq.view(bs, seqlen, num_head, head_dim)
-    xk = xk.view(bs, seqlen + prompt_cache_len, num_head, head_dim)
-    xv = xv.view(bs, seqlen + prompt_cache_len, num_head, head_dim)
-    mask_cache = torch.ones((seqlen, prompt_cache_len)).cuda().unsqueeze(0).unsqueeze(0).cuda()
-    mask = torch.tril(torch.ones(seqlen, seqlen), diagonal=0).unsqueeze(0).unsqueeze(0).cuda()
-    mask[mask == 0.0] = -100000000.0
-    mask = torch.cat([mask_cache, mask], dim=-1)
-    mask = mask.repeat(bs, num_head, 1, 1)
-    keys = xk
-    values = xv
-    xq = xq.transpose(1, 2)
-    keys = keys.transpose(1, 2)
-    values = values.transpose(1, 2)
-    scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(head_dim)
-    scores = F.softmax(scores.float() + mask, dim=-1).type_as(xq)
-    output = torch.matmul(scores, values).transpose(1, 2).contiguous().reshape(-1, num_head, head_dim)
-    return output
-
-
-def test():
-    import torch
-    import numpy as np
-
-    Z, H, N_CTX, D_HEAD = 10, 6, 500, 96
-    dtype = torch.float16
-    Z = 1
-    q = torch.empty((Z * N_CTX, H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.1, std=0.2)
-    k = torch.empty((Z * N_CTX + 7000, H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.4, std=0.2)
-    v = torch.empty((Z * N_CTX + 7000, H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.3, std=0.2)
-    o = torch.empty((Z * N_CTX, H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.3, std=0.2)
-    req_to_token_indexs = torch.zeros((10, Z * N_CTX + 7000), dtype=torch.int32, device="cuda")
-    max_input_len = N_CTX
-    Z = 1
-    b_start_loc = torch.zeros((Z,), dtype=torch.int32, device="cuda")
-    b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda")
-    b_req_idx = torch.ones((Z,), dtype=torch.int32, device="cuda")
-    b_prompt_cache_len = torch.zeros(1, dtype=torch.int32, device="cuda")
-    b_prompt_cache_len[0] = 0
-    prompt_cache_len = 0
-
-    b_seq_len[0] = 500
-    b_req_idx[0] = 0
-    req_to_token_indexs[0][: prompt_cache_len + N_CTX] = torch.tensor(
-        np.arange(prompt_cache_len + N_CTX), dtype=torch.int32
-    ).cuda()
-
-    torch_out = []
-    start = 0
-    for i in range(Z):
-        end = start + b_seq_len[i]
-        torch_o = torch_att(
-            q[start:end],
-            k[start : end + prompt_cache_len],
-            v[start : end + prompt_cache_len],
-            1,
-            b_seq_len[i],
-            H,
-            D_HEAD,
-            prompt_cache_len,
-        )
-        start = end
-        torch_out.append(torch_o)
-
-    torch_out = torch.cat(torch_out, dim=0)
-
-    context_attention_fwd(
-        q,
-        k,
-        v,
-        o,
-        b_req_idx,
-        b_start_loc,
-        b_seq_len + prompt_cache_len,
-        b_prompt_cache_len,
-        max_input_len,
-        req_to_token_indexs,
-    )
-
-    # context_attention_fwd_no_prompt_cache(
-    #         q, k, v, o, b_start_loc, b_seq_len, max_input_len
-    #     )
-
-    print("max ", torch.max(torch.abs(torch_out - o)))
-    print("mean ", torch.mean(torch.abs(torch_out - o)))
-    assert torch.allclose(torch_out, o, atol=1e-2, rtol=0)
diff --git a/lightllm/models/phi3/triton_kernel/destindex_copy_kv.py b/lightllm/models/phi3/triton_kernel/destindex_copy_kv.py
deleted file mode 100644
index 4f31895ae0..0000000000
--- a/lightllm/models/phi3/triton_kernel/destindex_copy_kv.py
+++ /dev/null
@@ -1,192 +0,0 @@
-import torch
-
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def _fwd_kernel_destindex_copy_kv(
-    K,
-    Dest_loc,
-    Out,
-    stride_k_bs,
-    stride_k_h,
-    stride_k_d,
-    stride_o_bs,
-    stride_o_h,
-    stride_o_d,
-    head_num,
-    head_dim,
-    BLOCK_DMODEL: tl.constexpr,
-    BLOCK_HEAD: tl.constexpr,
-):
-    cur_index = tl.program_id(0)
-    offs_h = tl.arange(0, BLOCK_HEAD)
-    offs_d = tl.arange(0, BLOCK_DMODEL)
-
-    dest_index = tl.load(Dest_loc + cur_index)
-
-    k_ptrs = K + cur_index * stride_k_bs + stride_k_h * offs_h[:, None] + stride_k_d * offs_d[None, :]
-    o_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :]
-
-    k = tl.load(k_ptrs, mask=(offs_h[:, None] < head_num) & (offs_d[None, :] < head_dim), other=0.0)
-    tl.store(o_ptrs, k, mask=(offs_h[:, None] < head_num) & (offs_d[None, :] < head_dim))
-    return
-
-
-@torch.no_grad()
-def destindex_copy_kv(K, DestLoc, Out):
-    seq_len = DestLoc.shape[0]
-    head_num = K.shape[1]
-    head_dim = K.shape[2]
-    assert K.shape[1] == Out.shape[1] and K.shape[2] == Out.shape[2]
-    BLOCK_HEAD = triton.next_power_of_2(head_num)
-    BLOCK_DMODEL = triton.next_power_of_2(head_dim)
-    grid = (seq_len,)
-    num_warps = 1
-
-    _fwd_kernel_destindex_copy_kv[grid](
-        K,
-        DestLoc,
-        Out,
-        K.stride(0),
-        K.stride(1),
-        K.stride(2),
-        Out.stride(0),
-        Out.stride(1),
-        Out.stride(2),
-        head_num,
-        head_dim,
-        BLOCK_DMODEL=BLOCK_DMODEL,
-        BLOCK_HEAD=BLOCK_HEAD,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
-
-
-@triton.jit
-def _fwd_kernel_destindex_copy_quantize_kv(
-    K,
-    Dest_loc,
-    Out,
-    Out_scale,
-    stride_k_bs,
-    stride_k_h,
-    stride_k_d,
-    stride_o_bs,
-    stride_o_h,
-    stride_o_d,
-    stride_os_bs,
-    stride_os_h,
-    stride_os_d,
-    head_num,
-    head_dim,
-    BLOCK_DMODEL: tl.constexpr,
-    BLOCK_HEAD: tl.constexpr,
-):
-    cur_index = tl.program_id(0)
-    offs_h = tl.arange(0, BLOCK_HEAD)
-    offs_d = tl.arange(0, BLOCK_DMODEL)
-
-    dest_index = tl.load(Dest_loc + cur_index)
-    src_data = tl.load(
-        K + cur_index * stride_k_bs + offs_h[:, None] * stride_k_h + stride_k_d * offs_d[None, :],
-        mask=(offs_h[:, None] < head_num) & (offs_d[None, :] < head_dim),
-        other=0.0,
-    )
-    abs_data = tl.abs(src_data)
-    data_scale = (tl.max(abs_data, axis=1) / 127.0).to(Out_scale.dtype.element_ty)[:, None]
-    q_src_data = (src_data / data_scale).to(tl.int8)
-    o_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :]
-    os_ptrs = Out_scale + dest_index * stride_os_bs + stride_os_h * offs_h[:, None]
-    tl.store(o_ptrs, q_src_data, mask=(offs_h[:, None] < head_num) & (offs_d[None, :] < head_dim))
-    tl.store(os_ptrs, data_scale, mask=(offs_h[:, None] < head_num))
-
-
-@torch.no_grad()
-def destindex_copy_quantize_kv(K, DestLoc, Out, Out_scale):
-    seq_len = DestLoc.shape[0]
-    head_num = K.shape[1]
-    head_dim = K.shape[2]
-    assert K.shape[1] == Out.shape[1] and K.shape[2] == Out.shape[2]
-    BLOCK_HEAD = triton.next_power_of_2(head_num)
-    BLOCK_DMODEL = triton.next_power_of_2(head_dim)
-    grid = (seq_len,)
-    num_warps = 1
-
-    _fwd_kernel_destindex_copy_quantize_kv[grid](
-        K,
-        DestLoc,
-        Out,
-        Out_scale,
-        K.stride(0),
-        K.stride(1),
-        K.stride(2),
-        Out.stride(0),
-        Out.stride(1),
-        Out.stride(2),
-        Out_scale.stride(0),
-        Out_scale.stride(1),
-        Out_scale.stride(2),
-        head_num,
-        head_dim,
-        BLOCK_DMODEL=BLOCK_DMODEL,
-        BLOCK_HEAD=BLOCK_HEAD,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-    return
-
-
-def test1():
-    import time
-
-    B, N_CTX, H, D = 32, 1024, 12, 96
-    dest = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda()
-    src = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda()
-    dest_loc = torch.arange(0, B * N_CTX, dtype=torch.int32, device="cuda")
-
-    for _ in range(10):
-        destindex_copy_kv(src, dest_loc, dest)
-    torch.cuda.synchronize()
-    t1 = time.time()
-    for _ in range(1000):
-        destindex_copy_kv(src, dest_loc, dest)
-    torch.cuda.synchronize()
-    t2 = time.time()
-
-    print("Time cost ", t2 - t1)
-    print("max ", torch.max(torch.abs(dest - src)))
-    print("mean ", torch.mean(torch.abs(dest - src)))
-    assert torch.allclose(src, dest, atol=1e-2, rtol=0)
-
-
-def test2():
-    import time
-
-    B, N_CTX, H, D = 32, 1024, 12, 96
-    src = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda()
-    dest_loc = torch.arange(0, B * N_CTX, dtype=torch.int32).cuda()
-    value_dest = torch.randn((B * N_CTX, H, D), dtype=torch.float16).cuda().to(torch.int8)
-    scale_dest = torch.randn((B * N_CTX, H, 1), dtype=torch.float16).cuda()
-
-    for _ in range(10):
-        destindex_copy_quantize_kv(src, dest_loc, value_dest, scale_dest)
-    torch.cuda.synchronize()
-    t1 = time.time()
-    for _ in range(1000):
-        destindex_copy_quantize_kv(src, dest_loc, value_dest, scale_dest)
-    torch.cuda.synchronize()
-    t2 = time.time()
-
-    print("Time cost ", t2 - t1)
-    print("max ", torch.max(torch.abs(value_dest * scale_dest - src)))
-    print("mean ", torch.mean(torch.abs(value_dest * scale_dest - src)))
-    cos = torch.nn.CosineSimilarity(0)
-    print("cos ", cos(src.flatten().to(torch.float32), (value_dest * scale_dest).flatten().to(torch.float32)))
-
-
-if __name__ == "__main__":
-    test1()
-    test2()
diff --git a/lightllm/models/qwen/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen/layer_infer/transformer_layer_infer.py
index 7a4b2ca816..333870eb9d 100755
--- a/lightllm/models/qwen/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/qwen/layer_infer/transformer_layer_infer.py
@@ -1,7 +1,4 @@
 import torch
-import torch.functional as F
-import torch.distributed as dist
-import numpy as np
 from typing import Tuple
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
 from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd
@@ -12,8 +9,8 @@
 class QwenTransformerLayerInfer(LlamaTransformerLayerInfer):
     """ """
 
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         return
 
     def _get_qkv(self, input_emb, infer_state: QwenInferStateInfo, layer_weight: QwenTransformerLayerWeight):
diff --git a/lightllm/models/qwen/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/qwen/layer_weights/pre_and_post_layer_weight.py
index 00f68eee69..52d1a54f59 100644
--- a/lightllm/models/qwen/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/qwen/layer_weights/pre_and_post_layer_weight.py
@@ -1,21 +1,26 @@
-import torch
-import numpy as np
 from lightllm.common.basemodel import PreAndPostLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, NoTpNormWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, RMSNormWeight
 
 
 class QwenPreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
+        hidden_size = network_config["hidden_size"]
+        vocab_size = network_config["vocab_size"]
         self.wte_weight_ = EmbeddingWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
             weight_name="transformer.wte.weight",
             data_type=self.data_type_,
         )
         self.lm_head_weight_ = LMHeadWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
             weight_name="lm_head.weight",
             data_type=self.data_type_,
         )
-        self.final_norm_weight_ = NoTpNormWeight(
+        self.final_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="transformer.ln_f.weight",
             data_type=self.data_type_,
         )
diff --git a/lightllm/models/qwen/layer_weights/transformer_layer_weight.py b/lightllm/models/qwen/layer_weights/transformer_layer_weight.py
index 9afb964ad2..ac1bf91f4b 100755
--- a/lightllm/models/qwen/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/qwen/layer_weights/transformer_layer_weight.py
@@ -3,8 +3,8 @@
 
 
 class QwenTransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
 
     def load_hf_weights(self, weights):
         qkv_weight_name = f"transformer.h.{self.layer_num_}.attn.c_attn.weight"
diff --git a/lightllm/models/qwen2/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/qwen2/layer_weights/pre_and_post_layer_weight.py
index a8a57c02ed..6449430d9e 100644
--- a/lightllm/models/qwen2/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/qwen2/layer_weights/pre_and_post_layer_weight.py
@@ -2,6 +2,6 @@
 
 
 class Qwen2PreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
         return
diff --git a/lightllm/models/qwen2/layer_weights/transformer_layer_weight.py b/lightllm/models/qwen2/layer_weights/transformer_layer_weight.py
index 6962818c49..fe6a5a2d49 100644
--- a/lightllm/models/qwen2/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/qwen2/layer_weights/transformer_layer_weight.py
@@ -2,50 +2,11 @@
 
 
 class Qwen2TransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
 
     def _init_weight_names(self):
         super()._init_weight_names()
         self._q_bias_name = f"model.layers.{self.layer_num_}.self_attn.q_proj.bias"
         self._k_bias_name = f"model.layers.{self.layer_num_}.self_attn.k_proj.bias"
         self._v_bias_name = f"model.layers.{self.layer_num_}.self_attn.v_proj.bias"
-
-    def _parse_config(self):
-        self.tp_q_head_num_ = self.network_config_["num_attention_heads"] // self.tp_world_size_
-        self.tp_k_head_num_ = max(self.network_config_["num_key_value_heads"] // self.tp_world_size_, 1)
-        self.tp_v_head_num_ = self.tp_k_head_num_
-        self.tp_o_head_num_ = self.tp_q_head_num_
-        head_dim = self.network_config_["hidden_size"] // self.network_config_["num_attention_heads"]
-        self.head_dim = self.network_config_.get("head_dim", head_dim)
-        assert (self.tp_k_head_num_ * self.tp_world_size_) % self.network_config_["num_key_value_heads"] == 0
-
-    def _repeat_weight(self, name, weights):
-        # for tp_world_size_ > num_key_value_heads
-        if name not in weights:
-            return
-
-        tensor = weights[name]
-        num_kv_heads = self.network_config_["num_key_value_heads"]
-        repeat_size = (self.tp_k_head_num_ * self.tp_world_size_) // num_kv_heads
-
-        if tensor.ndim == 1:
-            # Bias (1D tensor)
-            tensor = tensor.reshape(num_kv_heads, -1).unsqueeze(1).repeat(1, repeat_size, 1).reshape(-1)
-        else:
-            # Weight (2D tensor)
-            tensor = (
-                tensor.reshape(num_kv_heads, -1, tensor.shape[-1])
-                .unsqueeze(1)
-                .repeat(1, repeat_size, 1, 1)
-                .reshape(-1, tensor.shape[-1])
-            )
-        weights[name] = tensor
-
-    def load_hf_weights(self, weights):
-        self._repeat_weight(self._k_weight_name, weights)
-        self._repeat_weight(self._v_weight_name, weights)
-        if self._k_bias_name is not None and self._v_bias_name is not None:
-            self._repeat_weight(self._k_bias_name, weights)
-            self._repeat_weight(self._v_bias_name, weights)
-        return super().load_hf_weights(weights)
diff --git a/lightllm/models/qwen2/model.py b/lightllm/models/qwen2/model.py
index d2f067c42c..106610ff09 100644
--- a/lightllm/models/qwen2/model.py
+++ b/lightllm/models/qwen2/model.py
@@ -18,7 +18,7 @@ def __init__(self, kvargs):
 
     def _init_config(self):
         super()._init_config()
-        if self.config["sliding_window"] is None:
+        if self.config.get("sliding_window", None) is None:
             self.config["sliding_window"] = self.max_total_token_num
         # rename key [SYM: to be confirmed]
         return
diff --git a/lightllm/models/qwen2_reward/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/qwen2_reward/layer_weights/pre_and_post_layer_weight.py
index 7cf6366223..3c974691d5 100644
--- a/lightllm/models/qwen2_reward/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/qwen2_reward/layer_weights/pre_and_post_layer_weight.py
@@ -5,8 +5,8 @@
 
 
 class Qwen2RewardPreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
         del self.lm_head_weight_
         self.score_up_weight_ = ROWMMWeight(
             weight_names="score.0.weight",
diff --git a/lightllm/models/qwen2_vl/infer_struct.py b/lightllm/models/qwen2_vl/infer_struct.py
index 838590325c..747be932d9 100644
--- a/lightllm/models/qwen2_vl/infer_struct.py
+++ b/lightllm/models/qwen2_vl/infer_struct.py
@@ -4,13 +4,10 @@
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
 from lightllm.common.basemodel.infer_struct import InferStateInfo
 from lightllm.models.qwen2_vl.triton_kernel.get_mrope_position_ids import get_mrope_position_triton
-from lightllm.models.llama.flashattention_infer_struct import FlashAttentionStateInfo
 from lightllm.utils.envs_utils import get_env_start_args
 
 
 class Qwen2VLInferStateInfo(LlamaInferStateInfo):
-    init_flash_attention_state_func = FlashAttentionStateInfo._init_flash_attention_state
-
     def __init__(self):
         super().__init__()
         self.position_cos = None
@@ -35,10 +32,6 @@ def init_some_extra_state(self, model):
         self.position_ids = self.position_ids.contiguous()
         self.position_cos = model._cos_cached[self.position_ids]
         self.position_sin = model._sin_cached[self.position_ids]
-        if get_env_start_args().enable_fa3:
-            self.max_seq_len = self.max_kv_seq_len
-            self.q_max_seq_len = self.max_q_seq_len
-            self.init_flash_attention_state_func(model)
         return
 
     def get_mrope_position(self, multimodal_params: List[dict]) -> torch.Tensor:
@@ -85,6 +78,6 @@ def get_mrope_position(self, multimodal_params: List[dict]) -> torch.Tensor:
             position_ids=position_ids,
             b_ready_cache_len=self.b_ready_cache_len,
             b_q_seq_len=self.b_q_seq_len,
-            b_start_loc=self.b_start_loc,
+            b_start_loc=self.b_q_start_loc,
         )
         return position_ids
diff --git a/lightllm/models/qwen2_vl/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen2_vl/layer_infer/transformer_layer_infer.py
index 19e17c36e8..298a77044c 100755
--- a/lightllm/models/qwen2_vl/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/qwen2_vl/layer_infer/transformer_layer_infer.py
@@ -5,8 +5,8 @@
 
 
 class Qwen2VLTransformerLayerInfer(LlamaTransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         mrope_section = network_config["rope_scaling"]["mrope_section"]
         self.mrope_section = torch.tensor(mrope_section, dtype=torch.int32, device="cuda")
 
diff --git a/lightllm/models/qwen2_vl/model.py b/lightllm/models/qwen2_vl/model.py
index 61dd06773f..237c4ad897 100644
--- a/lightllm/models/qwen2_vl/model.py
+++ b/lightllm/models/qwen2_vl/model.py
@@ -79,6 +79,10 @@ def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs):
                     raise ValueError("image token error")
             except ValueError:
                 break
+        if multimodal_params:
+            image_cnt = len(multimodal_params.images)
+            if image_cnt != image_id:
+                raise ValueError(image_cnt == image_id, f"invalid image tag num: {image_cnt} vs {image_id}!")
         input_ids.extend(origin_ids)
         return input_ids
 
@@ -95,9 +99,6 @@ def __init__(self, kvargs):
         super().__init__(kvargs)
         return
 
-    def _init_inferstate_cls(self):
-        pass
-
     def _init_config(self):
         with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file:
             self.config = json.load(json_file)
diff --git a/lightllm/models/qwen2_vl/vision_process.py b/lightllm/models/qwen2_vl/vision_process.py
index 1c4f60794d..f2cd38ec8e 100644
--- a/lightllm/models/qwen2_vl/vision_process.py
+++ b/lightllm/models/qwen2_vl/vision_process.py
@@ -184,6 +184,8 @@ def preprocess(self, image) -> Tuple[torch.Tensor, torch.Tensor]:
             return self._preprocess_bydevice(image, device="cpu")
 
     def _preprocess_bydevice(self, image, device="cuda") -> Tuple[torch.Tensor, torch.Tensor]:
+        if image.mode != "RGB":
+            image = image.convert("RGB")
         image_arr = np.asarray(image, dtype=np.uint8)
         image_data = torch.from_numpy(image_arr).permute(2, 0, 1).contiguous().to(device=device, non_blocking=True)
 
diff --git a/lightllm/models/qwen3/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3/layer_infer/transformer_layer_infer.py
index 20f135e761..82331f8fb8 100644
--- a/lightllm/models/qwen3/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/qwen3/layer_infer/transformer_layer_infer.py
@@ -4,17 +4,14 @@
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
 from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd
-from lightllm.models.llama.triton_kernel.silu_and_mul import silu_and_mul_fwd
-from lightllm.models.qwen3.triton_kernel.qk_norm import qk_rmsnorm_forward
-from functools import partial
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
 
 
 class Qwen3TransformerLayerInfer(LlamaTransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         self.head_dim_ = network_config["head_dim"]
         return
 
@@ -27,14 +24,12 @@ def _get_qkv(
         input = input.view(-1, self.embed_dim_)
         q = layer_weight.q_proj.mm(input)
         cache_kv = layer_weight.kv_proj.mm(input)
-        qk_rmsnorm_forward(
+        layer_weight.q_norm_weight_(
             q,
-            weight=layer_weight.q_norm_weight_.weight,
             eps=self.eps_,
         )
-        qk_rmsnorm_forward(
+        layer_weight.k_norm_weight_(
             cache_kv[:, : self.tp_k_head_num_ * self.head_dim_],
-            weight=layer_weight.k_norm_weight_.weight,
             eps=self.eps_,
         )
         cache_kv = cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
diff --git a/lightllm/models/qwen3/layer_weights/transformer_layer_weight.py b/lightllm/models/qwen3/layer_weights/transformer_layer_weight.py
index 86b9e172a9..7d2163f283 100644
--- a/lightllm/models/qwen3/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/qwen3/layer_weights/transformer_layer_weight.py
@@ -1,12 +1,12 @@
 from lightllm.models.qwen2.layer_weights.transformer_layer_weight import Qwen2TransformerLayerWeight
 from lightllm.common.basemodel.layer_weights.meta_weights import (
-    NoTpNormWeight,
+    QKRMSNORMWeight,
 )
 
 
 class Qwen3TransformerLayerWeight(Qwen2TransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _init_weight_names(self):
@@ -19,6 +19,13 @@ def _init_weight_names(self):
 
     def _init_norm(self):
         super()._init_norm()
-
-        self.q_norm_weight_ = NoTpNormWeight(weight_name=self._q_norm_name, data_type=self.data_type_)
-        self.k_norm_weight_ = NoTpNormWeight(weight_name=self._k_norm_name, data_type=self.data_type_)
+        self.q_norm_weight_ = QKRMSNORMWeight(
+            dim=self.head_dim,
+            weight_name=self._q_norm_name,
+            data_type=self.data_type_,
+        )
+        self.k_norm_weight_ = QKRMSNORMWeight(
+            dim=self.head_dim,
+            weight_name=self._k_norm_name,
+            data_type=self.data_type_,
+        )
diff --git a/lightllm/models/qwen3/model.py b/lightllm/models/qwen3/model.py
index 21e71e0e02..e48b36e0f7 100644
--- a/lightllm/models/qwen3/model.py
+++ b/lightllm/models/qwen3/model.py
@@ -1,5 +1,3 @@
-import torch
-from typing import final
 from lightllm.models.registry import ModelRegistry
 from lightllm.models.qwen3.layer_infer.transformer_layer_infer import Qwen3TransformerLayerInfer
 from lightllm.models.qwen3.layer_weights.transformer_layer_weight import Qwen3TransformerLayerWeight
diff --git a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py
index 10a734e5c3..71b16cb34b 100644
--- a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py
@@ -14,12 +14,13 @@
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.dist_utils import get_global_world_size
 from lightllm.distributed.communication_op import all_gather_into_tensor, reduce_scatter_tensor
+from lightllm.utils.envs_utils import get_env_start_args
 
 logger = init_logger(__name__)
 
 
 class Qwen3MOETransformerLayerInfer(LlamaTransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
+    def __init__(self, layer_num, network_config):
         self.n_routed_experts = network_config["num_experts"]
         self.is_moe = (
             network_config["num_experts"] > 0
@@ -28,7 +29,7 @@ def __init__(self, layer_num, network_config, mode=[]):
         )
         self.num_experts_per_tok = network_config["num_experts_per_tok"]
         self.norm_topk_prob = network_config["norm_topk_prob"]
-        super().__init__(layer_num, network_config, mode)
+        super().__init__(layer_num, network_config)
         self.head_dim_ = network_config["head_dim"]
         self.tp_k_head_num_ = max(self.tp_k_head_num_, 1)
         self.tp_v_head_num_ = max(self.tp_v_head_num_, 1)
@@ -41,8 +42,8 @@ def _bind_func(self):
 
     def _bind_ffn(self):
         if self.is_moe:
-            moe_mode = os.environ.get("MOE_MODE", "TP")
-            if moe_mode == "EP":
+            enable_ep_moe = get_env_start_args().enable_ep_moe
+            if enable_ep_moe:
                 self._ffn = partial(Qwen3MOETransformerLayerInfer._moe_ffn_edp, self)
                 self._tpsp_ffn = self._tpsp_ffn_ep
             else:
@@ -60,20 +61,13 @@ def _get_qkv(
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         input = input.view(-1, self.embed_dim_)
         q = layer_weight.q_proj.mm(input)
-        cache_kv = layer_weight.kv_proj.mm(input).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
-
-        layer_weight.q_norm_weight_.rmsnorm_forward(
-            q.view(-1, self.head_dim_),
+        cache_kv = layer_weight.kv_proj.mm(input)
+        layer_weight.q_norm_weight_(q, eps=self.eps_)
+        layer_weight.k_norm_weight_(
+            cache_kv[:, : self.tp_k_head_num_ * self.head_dim_],
             eps=self.eps_,
-            out=q.view(-1, self.head_dim_),
         )
-
-        cache_kv[:, : self.tp_k_head_num_, :] = layer_weight.k_norm_weight_.rmsnorm_forward(
-            input=cache_kv[:, : self.tp_k_head_num_, :].reshape(-1, cache_kv.shape[-1]),
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
-        ).view(-1, self.tp_k_head_num_, cache_kv.shape[-1])
-
+        cache_kv = cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
         rotary_emb_fwd(
             q.view(-1, self.tp_q_head_num_, self.head_dim_),
             cache_kv[:, : self.tp_k_head_num_, :],
@@ -98,19 +92,13 @@ def _tpsp_get_qkv(
 
         input = input.view(-1, self.embed_dim_)
         q = layer_weight.q_proj.mm(input)
-        cache_kv = layer_weight.kv_proj.mm(input).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
-
-        layer_weight.q_norm_weight_.rmsnorm_forward(
-            q.view(-1, self.head_dim_),
+        cache_kv = layer_weight.kv_proj.mm(input)
+        layer_weight.q_norm_weight_(q, eps=self.eps_)
+        layer_weight.k_norm_weight_(
+            cache_kv[:, : self.tp_k_head_num_ * self.head_dim_],
             eps=self.eps_,
-            out=q.view(-1, self.head_dim_),
         )
-
-        cache_kv[:, : self.tp_k_head_num_, :] = layer_weight.k_norm_weight_.rmsnorm_forward(
-            cache_kv[:, : self.tp_k_head_num_, :].reshape(-1, cache_kv.shape[-1]),
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
-        ).view(-1, self.tp_k_head_num_, cache_kv.shape[-1])
+        cache_kv = cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
 
         rotary_emb_fwd(
             q.view(-1, self.tp_q_head_num_, self.head_dim_),
diff --git a/lightllm/models/qwen3_moe/layer_weights/transformer_layer_weight.py b/lightllm/models/qwen3_moe/layer_weights/transformer_layer_weight.py
index 72721f9d6f..a889609d7f 100644
--- a/lightllm/models/qwen3_moe/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/qwen3_moe/layer_weights/transformer_layer_weight.py
@@ -1,17 +1,17 @@
 import os
 from lightllm.models.qwen3.layer_weights.transformer_layer_weight import Qwen3TransformerLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import ROWMMWeight, FusedMoeWeightEP, create_tp_moe_wegiht_obj
+from lightllm.common.basemodel.layer_weights.meta_weights import ROWMMWeight, FusedMoeWeight
 
 
 class Qwen3MOETransformerLayerWeight(Qwen3TransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
         self.n_routed_experts = network_config["num_experts"]
         self.is_moe = (
             network_config["num_experts"] > 0
             and layer_num not in network_config["mlp_only_layers"]
             and (layer_num + 1) % network_config["decoder_sparse_step"] == 0
         )
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _init_weight_names(self):
@@ -32,15 +32,6 @@ def _init_weight_names(self):
         self._ffn_norm_weight_name = f"model.layers.{self.layer_num_}.post_attention_layernorm.weight"
         self._ffn_norm_bias_name = None
 
-    def load_hf_weights(self, weights):
-        kv_b_quant_method = self.quant_cfg.get_quant_method(self.layer_num_, "kv_b_proj")
-        if self.quant_cfg.quantized_weight:
-            _k_scale_weight_name = self._k_weight_name.replace("weight", kv_b_quant_method.weight_scale_suffix)
-            self._repeat_weight(_k_scale_weight_name, weights)
-            _v_scale_weight_name = self._v_weight_name.replace("weight", kv_b_quant_method.weight_scale_suffix)
-            self._repeat_weight(_v_scale_weight_name, weights)
-        return super().load_hf_weights(weights)
-
     def _init_weight(self):
         self._init_qkv()
         self._init_o()
@@ -53,42 +44,25 @@ def _init_weight(self):
     def _init_moe(self):
         moe_intermediate_size = self.network_config_["moe_intermediate_size"]
         self.moe_gate = ROWMMWeight(
+            in_dim=self.network_config_["hidden_size"],
+            out_dims=[self.n_routed_experts],
             weight_names=f"model.layers.{self.layer_num_}.mlp.gate.weight",
             data_type=self.data_type_,
-            layer_num=self.layer_num_,
-            name="moe_gate",
+            quant_method=None,
             tp_rank=0,
             tp_world_size=1,
         )
-        moe_mode = os.getenv("MOE_MODE", "TP")
-        assert moe_mode in ["EP", "TP"]
-        if moe_mode == "TP":
-            self.experts = create_tp_moe_wegiht_obj(
-                gate_proj_name="gate_proj",
-                down_proj_name="down_proj",
-                up_proj_name="up_proj",
-                e_score_correction_bias_name="",
-                weight_prefix=f"model.layers.{self.layer_num_}.mlp.experts",
-                n_routed_experts=self.n_routed_experts,
-                split_inter_size=moe_intermediate_size // self.tp_world_size_,
-                data_type=self.data_type_,
-                network_config=self.network_config_,
-                layer_num=self.layer_num_,
-                quant_cfg=self.quant_cfg,
-                num_fused_shared_experts=0,
-            )
-        elif moe_mode == "EP":
-            self.experts = FusedMoeWeightEP(
-                gate_proj_name="gate_proj",
-                down_proj_name="down_proj",
-                up_proj_name="up_proj",
-                e_score_correction_bias_name="",
-                weight_prefix=f"model.layers.{self.layer_num_}.mlp.experts",
-                n_routed_experts=self.n_routed_experts,
-                data_type=self.data_type_,
-                network_config=self.network_config_,
-                layer_num=self.layer_num_,
-                quant_cfg=self.quant_cfg,
-            )
-        else:
-            raise ValueError(f"Unsupported moe mode: {moe_mode}")
+        self.experts = FusedMoeWeight(
+            gate_proj_name="gate_proj",
+            down_proj_name="down_proj",
+            up_proj_name="up_proj",
+            e_score_correction_bias_name="",
+            weight_prefix=f"model.layers.{self.layer_num_}.mlp.experts",
+            n_routed_experts=self.n_routed_experts,
+            hidden_size=self.network_config_["hidden_size"],
+            moe_intermediate_size=moe_intermediate_size,
+            data_type=self.data_type_,
+            quant_method=self.quant_cfg.get_quant_method(self.layer_num_, "fused_moe"),
+            layer_num=self.layer_num_,
+            network_config=self.network_config_,
+        )
diff --git a/lightllm/models/qwen3_moe_mtp/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_moe_mtp/layer_infer/transformer_layer_infer.py
index d219173401..4e2b65d743 100644
--- a/lightllm/models/qwen3_moe_mtp/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/qwen3_moe_mtp/layer_infer/transformer_layer_infer.py
@@ -14,8 +14,8 @@
 
 
 class Qwen3MOEMTPTransformerLayerInfer(Qwen3MOETransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         return
 
     def context_forward(self, input_embdings, infer_state: InferStateInfo, layer_weight):
diff --git a/lightllm/models/qwen3_moe_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/qwen3_moe_mtp/layer_weights/pre_and_post_layer_weight.py
index 6cc447a594..e3a557d55c 100644
--- a/lightllm/models/qwen3_moe_mtp/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/qwen3_moe_mtp/layer_weights/pre_and_post_layer_weight.py
@@ -4,14 +4,15 @@
     EmbeddingWeight,
     ROWMMWeight,
     LMHeadWeight,
-    NoTpNormWeight,
+    RMSNormWeight,
 )
 
 
 class Qwen3MOEMTPPreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
 
+        hidden_size = network_config["hidden_size"]
         self.eh_proj_weight_ = ROWMMWeight(
             weight_names="model.layers.0.proj.weight",
             data_type=self.data_type_,
@@ -19,12 +20,14 @@ def __init__(self, data_type, network_config, mode):
             tp_rank=0,
             tp_world_size=1,
         )
-        self.enorm_weight_ = NoTpNormWeight(
+        self.enorm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="model.layers.0.norm_after_embedding.weight",
             data_type=self.data_type_,
             bias_name=None,
         )
-        self.hnorm_weight_ = NoTpNormWeight(
+        self.hnorm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="model.layers.0.norm_before_output.weight",
             data_type=self.data_type_,
             bias_name=None,
@@ -32,5 +35,5 @@ def __init__(self, data_type, network_config, mode):
         # 与Qwen3MOE模型共享
         self.wte_weight_: EmbeddingWeight = None
         self.lm_head_weight_: LMHeadWeight = None
-        self.final_norm_weight_: NoTpNormWeight = None
+        self.final_norm_weight_: RMSNormWeight = None
         return
diff --git a/lightllm/models/qwen3_moe_mtp/layer_weights/transformer_layer_weight.py b/lightllm/models/qwen3_moe_mtp/layer_weights/transformer_layer_weight.py
index 22d4d19505..2a11724ced 100644
--- a/lightllm/models/qwen3_moe_mtp/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/qwen3_moe_mtp/layer_weights/transformer_layer_weight.py
@@ -1,11 +1,11 @@
 import os
 from lightllm.models.qwen3_moe.layer_weights.transformer_layer_weight import Qwen3MOETransformerLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import NoTpNormWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import RMSNormWeight
 
 
 class Qwen3MOEMTPTransformerLayerWeight(Qwen3MOETransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _init_weight(self):
@@ -16,6 +16,10 @@ def _init_weight(self):
             self._init_ffn()
 
     def _init_norm(self):
-        self.ffn_norm_weight_ = NoTpNormWeight(
-            self._ffn_norm_weight_name, self.data_type_, bias_name=self._ffn_norm_bias_name
+        hidden_size = self.network_config_["hidden_size"]
+        self.ffn_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
+            weight_name=self._ffn_norm_weight_name,
+            data_type=self.data_type_,
+            bias_name=self._ffn_norm_bias_name,
         )
diff --git a/lightllm/models/qwen3_vl/layer_infer/pre_layer_infer.py b/lightllm/models/qwen3_vl/layer_infer/pre_layer_infer.py
index 96e453ebe7..c24166e13d 100644
--- a/lightllm/models/qwen3_vl/layer_infer/pre_layer_infer.py
+++ b/lightllm/models/qwen3_vl/layer_infer/pre_layer_infer.py
@@ -8,8 +8,8 @@
 
 
 class Qwen3VLMultimodalPreLayerInfer(LlamaMultimodalPreLayerInfer):
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         return
 
     def context_forward(
diff --git a/lightllm/models/qwen3_vl/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_vl/layer_infer/transformer_layer_infer.py
index 175340a77b..d34babaabe 100644
--- a/lightllm/models/qwen3_vl/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/qwen3_vl/layer_infer/transformer_layer_infer.py
@@ -1,29 +1,21 @@
 import torch
-import torch.functional as F
 import torch.distributed as dist
-import numpy as np
-from functools import partial
 from typing import Tuple
 from lightllm.common.basemodel.infer_struct import InferStateInfo
 from lightllm.models.qwen2_vl.triton_kernel.mrope import mrope_triton_fused
-from lightllm.models.qwen3.layer_infer.transformer_layer_infer import Qwen3TransformerLayerInfer
 from lightllm.models.qwen3.layer_weights.transformer_layer_weight import Qwen3TransformerLayerWeight
-from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
 from lightllm.models.qwen3_vl.infer_struct import Qwen3VLInferStateInfo
-from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd
-from lightllm.models.llama.triton_kernel.silu_and_mul import silu_and_mul_fwd
 from lightllm.distributed import all_reduce
-from lightllm.utils.dist_utils import get_global_world_size
 from lightllm.models.qwen3_vl.triton_kernel.deepstack_multimodal_emb import apply_deepstack_features
 from lightllm.models.qwen2_vl.layer_infer.transformer_layer_infer import Qwen2VLTransformerLayerInfer
-from lightllm.models.qwen3.triton_kernel.qk_norm import qk_rmsnorm_forward
 from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor
 
 
 class Qwen3VLTransformerLayerInfer(Qwen2VLTransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
+        self.head_dim_ = network_config["head_dim"]
         self.mrope_section = torch.tensor(
             network_config["rope_scaling"]["mrope_section"], dtype=torch.int32, device="cuda"
         )
@@ -37,14 +29,12 @@ def _get_qkv(
         input = input.view(-1, self.embed_dim_)
         q = layer_weight.q_proj.mm(input)
         cache_kv = layer_weight.kv_proj.mm(input)
-        qk_rmsnorm_forward(
+        layer_weight.q_norm_weight_(
             q,
-            weight=layer_weight.q_norm_weight_.weight,
             eps=self.eps_,
         )
-        qk_rmsnorm_forward(
+        layer_weight.k_norm_weight_(
             cache_kv[:, : self.tp_k_head_num_ * self.head_dim_],
-            weight=layer_weight.k_norm_weight_.weight,
             eps=self.eps_,
         )
         cache_kv = cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
diff --git a/lightllm/models/qwen3_vl/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/qwen3_vl/layer_weights/pre_and_post_layer_weight.py
index 5d41d85515..8a380853de 100644
--- a/lightllm/models/qwen3_vl/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/qwen3_vl/layer_weights/pre_and_post_layer_weight.py
@@ -12,8 +12,8 @@ def rename_weight_keys(weights):
 
 
 class Qwen3VLPreAndPostLayerWeight(Qwen2PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
         return
 
     def load_hf_weights(self, weights):
diff --git a/lightllm/models/qwen3_vl/model.py b/lightllm/models/qwen3_vl/model.py
index 0d8a81f671..74aa33e3c0 100644
--- a/lightllm/models/qwen3_vl/model.py
+++ b/lightllm/models/qwen3_vl/model.py
@@ -37,9 +37,6 @@ def __init__(self, kvargs):
         super().__init__(kvargs)
         return
 
-    def _init_inferstate_cls(self):
-        pass
-
     def _init_config(self):
         with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file:
             all_config = json.load(json_file)
diff --git a/lightllm/models/qwen3_vl_moe/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_vl_moe/layer_infer/transformer_layer_infer.py
index b155f8b907..4ccc6da372 100644
--- a/lightllm/models/qwen3_vl_moe/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/qwen3_vl_moe/layer_infer/transformer_layer_infer.py
@@ -5,15 +5,13 @@
 from lightllm.models.qwen3_moe.layer_infer.transformer_layer_infer import Qwen3MOETransformerLayerInfer
 from lightllm.models.qwen3_moe.layer_weights.transformer_layer_weight import Qwen3MOETransformerLayerWeight
 from lightllm.models.qwen3_vl.infer_struct import Qwen3VLInferStateInfo
-from lightllm.models.qwen3.triton_kernel.qk_norm import qk_rmsnorm_forward
 from lightllm.distributed import all_reduce
-from lightllm.utils.dist_utils import get_global_world_size
 from lightllm.models.qwen3_vl.triton_kernel.deepstack_multimodal_emb import apply_deepstack_features
 
 
 class Qwen3VLMOETransformerLayerInfer(Qwen3MOETransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         self.mrope_section = torch.tensor(
             network_config["rope_scaling"]["mrope_section"], dtype=torch.int32, device="cuda"
         )
@@ -27,14 +25,12 @@ def _get_qkv(
         input = input.view(-1, self.embed_dim_)
         q = layer_weight.q_proj.mm(input)
         cache_kv = layer_weight.kv_proj.mm(input)
-        qk_rmsnorm_forward(
+        layer_weight.q_norm_weight_(
             q,
-            weight=layer_weight.q_norm_weight_.weight,
             eps=self.eps_,
         )
-        qk_rmsnorm_forward(
+        layer_weight.k_norm_weight_(
             cache_kv[:, : self.tp_k_head_num_ * self.head_dim_],
-            weight=layer_weight.k_norm_weight_.weight,
             eps=self.eps_,
         )
         cache_kv = cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
diff --git a/lightllm/models/qwen3_vl_moe/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/qwen3_vl_moe/layer_weights/pre_and_post_layer_weight.py
index b1f5ee6600..0ba06c6aed 100644
--- a/lightllm/models/qwen3_vl_moe/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/qwen3_vl_moe/layer_weights/pre_and_post_layer_weight.py
@@ -1,24 +1,28 @@
-import numpy as np
 from lightllm.common.basemodel import PreAndPostLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, NoTpNormWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, RMSNormWeight
 
 
 class Qwen3VLMOEPreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
+        hidden_size = network_config["hidden_size"]
+        vocab_size = network_config["vocab_size"]
         self.wte_weight_ = EmbeddingWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
             weight_name="model.language_model.embed_tokens.weight",
             data_type=self.data_type_,
         )
         tie_word_embeddings = self.network_config_.get("tie_word_embeddings", False)
-        if tie_word_embeddings:
-            self.lm_head_weight_: LMHeadWeight = self.wte_weight_
-        else:
-            self.lm_head_weight_ = LMHeadWeight(
-                weight_name="lm_head.weight",
-                data_type=self.data_type_,
-            )
-        self.final_norm_weight_ = NoTpNormWeight(
+        self.lm_head_weight_ = LMHeadWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
+            weight_name="lm_head.weight",
+            data_type=self.data_type_,
+            embedding_weight=self.wte_weight_ if tie_word_embeddings else None,
+        )
+        self.final_norm_weight_ = RMSNormWeight(
+            dim=hidden_size,
             weight_name="model.language_model.norm.weight",
             data_type=self.data_type_,
         )
diff --git a/lightllm/models/qwen3_vl_moe/layer_weights/transformers_layer_weight.py b/lightllm/models/qwen3_vl_moe/layer_weights/transformers_layer_weight.py
index f4eef6e698..83c05ba264 100644
--- a/lightllm/models/qwen3_vl_moe/layer_weights/transformers_layer_weight.py
+++ b/lightllm/models/qwen3_vl_moe/layer_weights/transformers_layer_weight.py
@@ -1,11 +1,10 @@
 import os
 from lightllm.models.qwen3_moe.layer_weights.transformer_layer_weight import Qwen3MOETransformerLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import ROWMMWeight, FusedMoeWeightEP, create_tp_moe_wegiht_obj
 
 
 class Qwen3VLMOETransformerLayerWeight(Qwen3MOETransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
 
     def load_hf_weights(self, weights):
         moe_prefix = f"model.layers.{self.layer_num_}.mlp.experts"
diff --git a/lightllm/models/qwen3_vl_moe/model.py b/lightllm/models/qwen3_vl_moe/model.py
index b11f22fdb7..cc1201de2c 100644
--- a/lightllm/models/qwen3_vl_moe/model.py
+++ b/lightllm/models/qwen3_vl_moe/model.py
@@ -25,9 +25,6 @@ def __init__(self, kvargs):
         super().__init__(kvargs)
         return
 
-    def _init_inferstate_cls(self):
-        pass
-
     def _init_config(self):
         with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file:
             all_config = json.load(json_file)
diff --git a/lightllm/models/qwen_vl/layer_infer/pre_layer_infer.py b/lightllm/models/qwen_vl/layer_infer/pre_layer_infer.py
index f439073077..939843a3eb 100644
--- a/lightllm/models/qwen_vl/layer_infer/pre_layer_infer.py
+++ b/lightllm/models/qwen_vl/layer_infer/pre_layer_infer.py
@@ -24,8 +24,8 @@
 
 
 class LlamaMultimodalPreLayerInfer(LlamaPreLayerInfer):
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         return
 
     def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight):
diff --git a/lightllm/models/qwen_vl/model.py b/lightllm/models/qwen_vl/model.py
index edebccf17f..0c6fa31f47 100644
--- a/lightllm/models/qwen_vl/model.py
+++ b/lightllm/models/qwen_vl/model.py
@@ -1,5 +1,3 @@
-import json
-import numpy as np
 import unicodedata
 from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer
 from lightllm.server.core.objs import SamplingParams
@@ -88,7 +86,8 @@ def encode(self, prompt, multimodal_params: MultimodalParams = None):
         input_ids.extend(origin_ids[end:])
         if multimodal_params:
             image_cnt = len(multimodal_params.images)
-            assert image_cnt == image_id, "invalid image tag num: {} vs {}!".format(image_cnt, image_id)
+            if image_cnt != image_id:
+                raise ValueError(image_cnt == image_id, f"invalid image tag num: {image_cnt} vs {image_id}!")
         return input_ids
 
 
diff --git a/lightllm/models/stablelm/layer_infer/transformer_layer_infer.py b/lightllm/models/stablelm/layer_infer/transformer_layer_infer.py
index 395ed4ba1a..55848ce66a 100755
--- a/lightllm/models/stablelm/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/stablelm/layer_infer/transformer_layer_infer.py
@@ -8,8 +8,8 @@
 
 
 class StablelmTransformerLayerInfer(LlamaTransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         self.partial_rotary_factor = self.network_config_.get("partial_rotary_factor", 1)
         return
 
@@ -53,17 +53,13 @@ def _tpsp_get_o(self, input, infer_state, layer_weight) -> Tuple[torch.Tensor, t
     def _att_norm(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: StablelmTransformerLayerWeight
     ) -> torch.Tensor:
-        return layer_weight.att_norm_weight_.layernorm_forward(
-            input=input.view(-1, self.embed_dim_),
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
+        return layer_weight.att_norm_weight_(
+            input=input.view(-1, self.embed_dim_), eps=self.eps_, alloc_func=self.alloc_tensor
         )
 
     def _ffn_norm(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: StablelmTransformerLayerWeight
     ) -> torch.Tensor:
-        return layer_weight.ffn_norm_weight_.layernorm_forward(
-            input=input.view(-1, self.embed_dim_),
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
+        return layer_weight.ffn_norm_weight_(
+            input=input.view(-1, self.embed_dim_), eps=self.eps_, alloc_func=self.alloc_tensor
         )
diff --git a/lightllm/models/stablelm/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/stablelm/layer_weights/pre_and_post_layer_weight.py
index 0ad3e07df5..885c7ead7d 100755
--- a/lightllm/models/stablelm/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/stablelm/layer_weights/pre_and_post_layer_weight.py
@@ -1,10 +1,13 @@
-from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight, NoTpNormWeight
+from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import LayerNormWeight
 
 
 class StableLMPreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
-        self.final_norm_weight_ = NoTpNormWeight(
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
+        hidden_size = network_config["hidden_size"]
+        self.final_norm_weight_ = LayerNormWeight(
+            dim=hidden_size,
             weight_name="model.norm.weight",
             data_type=self.data_type_,
             bias_name="model.norm.bias",
diff --git a/lightllm/models/stablelm/layer_weights/transformer_layer_weight.py b/lightllm/models/stablelm/layer_weights/transformer_layer_weight.py
index a1a73f6745..03ee50feb5 100755
--- a/lightllm/models/stablelm/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/stablelm/layer_weights/transformer_layer_weight.py
@@ -2,8 +2,8 @@
 
 
 class StablelmTransformerLayerWeight(Qwen2TransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _init_weight_names(self):
diff --git a/lightllm/models/stablelm/model.py b/lightllm/models/stablelm/model.py
index 2ed710fd4c..a3d295358f 100644
--- a/lightllm/models/stablelm/model.py
+++ b/lightllm/models/stablelm/model.py
@@ -1,6 +1,3 @@
-import os
-import json
-import torch
 from lightllm.models.registry import ModelRegistry
 from lightllm.models.stablelm.layer_infer.transformer_layer_infer import StablelmTransformerLayerInfer
 from lightllm.models.bloom.layer_infer.post_layer_infer import BloomPostLayerInfer
diff --git a/lightllm/models/starcoder/layer_infer/pre_layer_infer.py b/lightllm/models/starcoder/layer_infer/pre_layer_infer.py
index 52072a3487..b3cd083c30 100644
--- a/lightllm/models/starcoder/layer_infer/pre_layer_infer.py
+++ b/lightllm/models/starcoder/layer_infer/pre_layer_infer.py
@@ -9,29 +9,23 @@
 class StarcoderPreLayerInfer(PreLayerInfer):
     """ """
 
-    def __init__(self, network_config, mode):
-        super().__init__(network_config, mode)
+    def __init__(self, network_config):
+        super().__init__(network_config)
         self.layer_norm_eps_ = network_config["layer_norm_epsilon"]
 
     def context_forward(self, input_ids, infer_state: InferStateInfo, layer_weight: StarcoderPreAndPostLayerWeight):
-        input_embdings = layer_weight.wte_weight_.embedding(input_ids=input_ids, alloc_func=self.alloc_tensor)
+        input_embdings = layer_weight.wte_weight_(input_ids=input_ids, alloc_func=self.alloc_tensor)
         if self.tp_world_size_ > 1:
             all_reduce(input_embdings, group=infer_state.dist_group, op=dist.ReduceOp.SUM, async_op=False)
 
-        position_embeds = layer_weight.wpe_weight_.embedding(
-            input_ids=infer_state.position_ids,
-            alloc_func=self.alloc_tensor,
-        )
+        position_embeds = layer_weight.wpe_weight_(input_ids=infer_state.position_ids, alloc_func=self.alloc_tensor)
 
         return input_embdings.add_(position_embeds)
 
     def token_forward(self, input_ids, infer_state: InferStateInfo, layer_weight: StarcoderPreAndPostLayerWeight):
-        input_embdings = layer_weight.wte_weight_.embedding(input_ids=input_ids, alloc_func=self.alloc_tensor)
+        input_embdings = layer_weight.wte_weight_(input_ids=input_ids, alloc_func=self.alloc_tensor)
         if self.tp_world_size_ > 1:
             all_reduce(input_embdings, group=infer_state.dist_group, op=dist.ReduceOp.SUM, async_op=False)
 
-        position_embeds = layer_weight.wpe_weight_.embedding(
-            input_ids=infer_state.position_ids,
-            alloc_func=self.alloc_tensor,
-        )
+        position_embeds = layer_weight.wpe_weight_(input_ids=infer_state.position_ids, alloc_func=self.alloc_tensor)
         return input_embdings.add_(position_embeds)
diff --git a/lightllm/models/starcoder/layer_infer/transformer_layer_infer.py b/lightllm/models/starcoder/layer_infer/transformer_layer_infer.py
index 018816fcc6..074f3411a7 100644
--- a/lightllm/models/starcoder/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/starcoder/layer_infer/transformer_layer_infer.py
@@ -1,17 +1,19 @@
 from lightllm.models.bloom.layer_infer.transformer_layer_infer import BloomTransformerLayerInfer
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
+from functools import partial
 
 
 class StarcoderTransformerLayerInfer(BloomTransformerLayerInfer):
     """ """
 
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
         self.tp_k_head_num_ = 1
         self.tp_v_head_num_ = 1
         self._bind_func()
         return
 
     def _bind_func(self):
-        LlamaTransformerLayerInfer._bind_attention(self)
+        self._context_attention_kernel = partial(LlamaTransformerLayerInfer._context_attention_kernel, self)
+        self._token_attention_kernel = partial(LlamaTransformerLayerInfer._token_attention_kernel, self)
         return
diff --git a/lightllm/models/starcoder/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/starcoder/layer_weights/pre_and_post_layer_weight.py
index d5bdd79a7b..a258480f65 100644
--- a/lightllm/models/starcoder/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/starcoder/layer_weights/pre_and_post_layer_weight.py
@@ -1,31 +1,42 @@
 from lightllm.common.basemodel import PreAndPostLayerWeight
 from lightllm.common.basemodel.layer_weights.meta_weights import (
     EmbeddingWeight,
-    NoTpNormWeight,
+    LayerNormWeight,
     NoTpPosEmbeddingWeight,
     LMHeadWeight,
 )
 
 
 class StarcoderPreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
 
+    def _create_weight(self):
+        hidden_size = self.network_config["hidden_size"]
+        vocab_size = self.network_config["vocab_size"]
+        max_position_embeddings = self.network_config["max_position_embeddings"]
         self.wte_weight_ = EmbeddingWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
             weight_name="transformer.wte.weight",
             data_type=self.data_type_,
         )
         self.wpe_weight_ = NoTpPosEmbeddingWeight(
+            dim=hidden_size,
+            max_position_embeddings=max_position_embeddings,
             weight_name="transformer.wpe.weight",
             data_type=self.data_type_,
         )
 
-        self.final_norm_weight_ = NoTpNormWeight(
+        self.final_norm_weight_ = LayerNormWeight(
+            dim=hidden_size,
             weight_name="transformer.ln_f.weight",
             bias_name="transformer.ln_f.bias",
             data_type=self.data_type_,
         )
         self.lm_head_weight_ = LMHeadWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
             weight_name="lm_head.weight",
             data_type=self.data_type_,
         )
diff --git a/lightllm/models/starcoder/layer_weights/transformer_layer_weight.py b/lightllm/models/starcoder/layer_weights/transformer_layer_weight.py
index 2aa9dd9ef2..4adbe4f5e6 100644
--- a/lightllm/models/starcoder/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/starcoder/layer_weights/transformer_layer_weight.py
@@ -3,8 +3,8 @@
 
 
 class StarcoderTransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg, layer_prefix="transformer.h")
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg, layer_prefix="transformer.h")
         assert network_config["num_attention_heads"] % self.tp_world_size_ == 0
 
     def load_hf_weights(self, weights):
@@ -51,18 +51,18 @@ def _init_weight_names(self):
 
     def _init_ffn(self):
         self.gate_up_proj = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_inter],
             weight_names=self._gate_up_weight_name,
             data_type=self.data_type_,
             bias_names=self._gate_up_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="gate_up_proj",
+            quant_method=self.get_quant_method("gate_up_proj"),
         )
         self.down_proj = COLMMWeight(
+            in_dim=self.n_inter,
+            out_dims=[self.n_embed],
             weight_names=self._down_weight_name,
             data_type=self.data_type_,
             bias_names=self._down_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="down_proj",
+            quant_method=self.get_quant_method("down_proj"),
         )
diff --git a/lightllm/models/starcoder2/layer_infer/transformer_layer_infer.py b/lightllm/models/starcoder2/layer_infer/transformer_layer_infer.py
index 796a96bc4a..3e32682ecb 100644
--- a/lightllm/models/starcoder2/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/starcoder2/layer_infer/transformer_layer_infer.py
@@ -5,25 +5,21 @@
 
 
 class Starcoder2TransformerLayerInfer(LlamaTransformerLayerInfer):
-    def __init__(self, layer_num, network_config, mode=[]):
-        super().__init__(layer_num, network_config, mode)
+    def __init__(self, layer_num, network_config):
+        super().__init__(layer_num, network_config)
 
     def _att_norm(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: Starcoder2TransformerLayerWeight
     ) -> torch.Tensor:
-        return layer_weight.att_norm_weight_.layernorm_forward(
-            input=input.view(-1, self.embed_dim_),
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
+        return layer_weight.att_norm_weight_(
+            input=input.view(-1, self.embed_dim_), eps=self.eps_, alloc_func=self.alloc_tensor
         )
 
     def _ffn_norm(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: Starcoder2TransformerLayerWeight
     ) -> torch.Tensor:
-        return layer_weight.ffn_norm_weight_.layernorm_forward(
-            input=input.view(-1, self.embed_dim_),
-            eps=self.eps_,
-            alloc_func=self.alloc_tensor,
+        return layer_weight.ffn_norm_weight_(
+            input=input.view(-1, self.embed_dim_), eps=self.eps_, alloc_func=self.alloc_tensor
         )
 
     def _ffn(
diff --git a/lightllm/models/starcoder2/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/starcoder2/layer_weights/pre_and_post_layer_weight.py
index 28a26cb4b3..cc256c442d 100644
--- a/lightllm/models/starcoder2/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/starcoder2/layer_weights/pre_and_post_layer_weight.py
@@ -1,26 +1,29 @@
-import numpy as np
 from lightllm.common.basemodel import PreAndPostLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, NoTpNormWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight, LMHeadWeight, LayerNormWeight
 
 
 class Starcoder2PreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
-
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
+        hidden_size = network_config["hidden_size"]
+        vocab_size = network_config["vocab_size"]
         self.wte_weight_ = EmbeddingWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
             weight_name="model.embed_tokens.weight",
             data_type=self.data_type_,
         )
         tie_word_embeddings = self.network_config_.get("tie_word_embeddings", False)
-        if tie_word_embeddings:
-            self.lm_head_weight_: LMHeadWeight = self.wte_weight_
-        else:
-            self.lm_head_weight_ = LMHeadWeight(
-                weight_name="lm_head.weight",
-                data_type=self.data_type_,
-            )
+        self.lm_head_weight_ = LMHeadWeight(
+            dim=hidden_size,
+            vocab_size=vocab_size,
+            weight_name="lm_head.weight",
+            data_type=self.data_type_,
+            embedding_weight=self.wte_weight_ if tie_word_embeddings else None,
+        )
 
-        self.final_norm_weight_ = NoTpNormWeight(
+        self.final_norm_weight_ = LayerNormWeight(
+            dim=hidden_size,
             weight_name="model.norm.weight",
             data_type=self.data_type_,
             bias_name="model.norm.bias",
diff --git a/lightllm/models/starcoder2/layer_weights/transformer_layer_weight.py b/lightllm/models/starcoder2/layer_weights/transformer_layer_weight.py
index 6314fa0e57..7370c69530 100644
--- a/lightllm/models/starcoder2/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/starcoder2/layer_weights/transformer_layer_weight.py
@@ -3,8 +3,8 @@
 
 
 class Starcoder2TransformerLayerWeight(LlamaTransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _parse_config(self):
@@ -28,18 +28,18 @@ def _init_weight_names(self):
 
     def _init_ffn(self):
         self.up_proj = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_inter],
             weight_names=self._up_weight_name,
             data_type=self.data_type_,
             bias_names=self._up_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="up_proj",
+            quant_method=self.get_quant_method("up_proj"),
         )
         self.down_proj = COLMMWeight(
+            in_dim=self.n_inter,
+            out_dims=[self.n_embed],
             weight_names=self._down_weight_name,
             data_type=self.data_type_,
             bias_names=self._down_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="down_proj",
+            quant_method=self.get_quant_method("down_proj"),
         )
diff --git a/lightllm/models/tarsier2/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/tarsier2/layer_weights/pre_and_post_layer_weight.py
index b24fc0f0d1..44e18c2826 100644
--- a/lightllm/models/tarsier2/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/tarsier2/layer_weights/pre_and_post_layer_weight.py
@@ -17,8 +17,8 @@ def rename_weight_keys(weights):
 
 
 class Tarsier2Qwen2PreAndPostLayerWeight(Qwen2PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
         return
 
     def load_hf_weights(self, weights):
@@ -28,8 +28,8 @@ def load_hf_weights(self, weights):
 
 
 class Tarsier2LlamaPreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
         return
 
     def load_hf_weights(self, weights):
diff --git a/lightllm/models/tarsier2/model.py b/lightllm/models/tarsier2/model.py
index dad252b979..10a7f368c4 100644
--- a/lightllm/models/tarsier2/model.py
+++ b/lightllm/models/tarsier2/model.py
@@ -78,6 +78,10 @@ def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs):
                     raise ValueError("image token error")
             except ValueError:
                 break
+        if multimodal_params:
+            image_cnt = len(multimodal_params.images)
+            if image_cnt != image_id:
+                raise ValueError(image_cnt == image_id, f"invalid image tag num: {image_cnt} vs {image_id}!")
         input_ids.extend(origin_ids[start_idx:])
         return input_ids
 
diff --git a/lightllm/models/vit/layer_infer/post_layer_infer.py b/lightllm/models/vit/layer_infer/post_layer_infer.py
index 613aec3fa7..fa4a87f158 100644
--- a/lightllm/models/vit/layer_infer/post_layer_infer.py
+++ b/lightllm/models/vit/layer_infer/post_layer_infer.py
@@ -9,11 +9,10 @@
 class ViTPostLayerInfer:
     """ """
 
-    def __init__(self, network_config, mode):
+    def __init__(self, network_config):
         self.tp_rank_ = get_current_rank_in_dp()
         self.tp_world_size_ = get_dp_world_size()
         self.network_config_ = network_config
-        self.mode = mode
         self.llm_hidden_size = network_config["llm_hidden_size"]
         self.downsample_ratio = network_config["downsample_ratio"]
         return
diff --git a/lightllm/models/vit/layer_infer/pre_layer_infer.py b/lightllm/models/vit/layer_infer/pre_layer_infer.py
index 896e8e898c..306bf9f0e6 100644
--- a/lightllm/models/vit/layer_infer/pre_layer_infer.py
+++ b/lightllm/models/vit/layer_infer/pre_layer_infer.py
@@ -11,11 +11,10 @@
 class ViTPreLayerInfer:
     """ """
 
-    def __init__(self, network_config, mode):
+    def __init__(self, network_config):
         self.tp_rank_ = get_current_rank_in_dp()
         self.tp_world_size_ = get_dp_world_size()
         self.network_config_ = network_config
-        self.mode = mode
         return
 
     def forward(self, pixel_values, layer_weight: ViTPreAndPostLayerWeight):
diff --git a/lightllm/models/vit/layer_infer/transformer_layer_infer.py b/lightllm/models/vit/layer_infer/transformer_layer_infer.py
index 0b89dca11e..0d55d1b57f 100644
--- a/lightllm/models/vit/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/vit/layer_infer/transformer_layer_infer.py
@@ -13,7 +13,7 @@
 class ViTTransformerLayerInfer:
     """ """
 
-    def __init__(self, layer_num, network_config, mode=[]):
+    def __init__(self, layer_num, network_config):
         self.tp_rank_ = get_current_rank_in_dp()
         self.tp_world_size_ = get_dp_world_size()
         self.eps_ = network_config["layer_norm_eps"]
@@ -25,7 +25,6 @@ def __init__(self, layer_num, network_config, mode=[]):
         self.tp_padding_embed_dim_ = self.tp_padding_head_num * self.head_dim_
 
         self.network_config_ = network_config
-        self.mode = mode
         self.layer_num_ = layer_num
         return
 
diff --git a/lightllm/models/vit/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/vit/layer_weights/pre_and_post_layer_weight.py
index 276d4e5d0b..79dc9d95cd 100644
--- a/lightllm/models/vit/layer_weights/pre_and_post_layer_weight.py
+++ b/lightllm/models/vit/layer_weights/pre_and_post_layer_weight.py
@@ -7,12 +7,44 @@
 
 
 class ViTPreAndPostLayerWeight(PreAndPostLayerWeight):
-    def __init__(self, data_type, network_config, mode):
-        super().__init__(data_type, network_config, mode)
+    def __init__(self, data_type, network_config):
+        super().__init__(data_type, network_config)
         self.embed_dim = self.network_config_["hidden_size"]
         self.image_size = self.network_config_["image_size"]
         self.patch_size = self.network_config_["patch_size"]
         self.llm_hidden_size = self.network_config_["llm_hidden_size"]
+        self._create_weight()
+        return
+
+    def _create_weight(self):
+        split_indexes = np.linspace(0, self.embed_dim, self.tp_world_size_ + 1, dtype=np.int64)
+        split_start = split_indexes[self.tp_rank_]
+        split_end = split_indexes[self.tp_rank_ + 1]
+        split_embed_dim = split_end - split_start
+
+        # Pre-allocate memory for vision model weights
+        self.class_embedding = torch.empty((1, 1, split_embed_dim), dtype=self.data_type_).cuda()
+        self.position_embedding = torch.empty(
+            (1, 197, split_embed_dim), dtype=self.data_type_
+        ).cuda()  # 197 = (224//16)^2 + 1
+        self.patch_embedding_weight_ = torch.empty(
+            (split_embed_dim, 3, self.patch_size, self.patch_size), dtype=self.data_type_
+        ).cuda()
+        self.patch_embedding_bias_ = torch.empty(split_embed_dim, dtype=self.data_type_).cuda()
+
+        # Pre-allocate memory for adapter weights
+        self.layernorm_weight_ = torch.empty(self.embed_dim, dtype=self.data_type_).cuda()
+        self.layernorm_bias_ = torch.empty(self.embed_dim, dtype=self.data_type_).cuda()
+
+        split_indexes_llm = np.linspace(0, self.llm_hidden_size, self.tp_world_size_ + 1, dtype=np.int64)
+        split_start_llm = split_indexes_llm[self.tp_rank_]
+        split_end_llm = split_indexes_llm[self.tp_rank_ + 1]
+        split_llm_hidden_size = split_end_llm - split_start_llm
+
+        self.mlp1_1_weight_ = torch.empty((self.llm_hidden_size, split_llm_hidden_size), dtype=self.data_type_).cuda()
+        self.mlp1_1_bias_ = torch.empty(split_llm_hidden_size, dtype=self.data_type_).cuda()
+        self.mlp1_3_weight_ = torch.empty((split_llm_hidden_size, self.llm_hidden_size), dtype=self.data_type_).cuda()
+        self.mlp1_3_bias_ = torch.empty(self.llm_hidden_size, dtype=self.data_type_).cuda()
         return
 
     def _cuda(self, cpu_tensor):
@@ -40,40 +72,38 @@ def load_hf_weights(self, weights):
         split_start = split_indexes[self.tp_rank_]
         split_end = split_indexes[self.tp_rank_ + 1]
         if "vision_model.embeddings.class_embedding" in weights:
-            self.class_embedding = self._cuda(
-                weights["vision_model.embeddings.class_embedding"][:, :, split_start:split_end]
-            )
+            self.class_embedding.copy_(weights["vision_model.embeddings.class_embedding"][:, :, split_start:split_end])
         if "vision_model.embeddings.position_embedding" in weights:
-            self.position_embedding = self._cuda(
+            self.position_embedding.copy_(
                 weights["vision_model.embeddings.position_embedding"][:, :, split_start:split_end]
             )
         if "vision_model.embeddings.patch_embedding.weight" in weights:
-            self.patch_embedding_weight_ = self._cuda(
+            self.patch_embedding_weight_.copy_(
                 weights["vision_model.embeddings.patch_embedding.weight"][split_start:split_end, :, :, :]
             )
         if "vision_model.embeddings.patch_embedding.bias" in weights:
-            self.patch_embedding_bias_ = self._cuda(
+            self.patch_embedding_bias_.copy_(
                 weights["vision_model.embeddings.patch_embedding.bias"][split_start:split_end]
             )
 
         if "mlp1.0.weight" in weights:
-            self.layernorm_weight_ = self._cuda(weights["mlp1.0.weight"])
+            self.layernorm_weight_.copy_(weights["mlp1.0.weight"])
         if "mlp1.0.bias" in weights:
-            self.layernorm_bias_ = self._cuda(weights["mlp1.0.bias"])
+            self.layernorm_bias_.copy_(weights["mlp1.0.bias"])
 
         split_indexes = np.linspace(0, self.llm_hidden_size, self.tp_world_size_ + 1, dtype=np.int64)
         split_start = split_indexes[self.tp_rank_]
         split_end = split_indexes[self.tp_rank_ + 1]
 
         if "mlp1.1.weight" in weights:
-            self.mlp1_1_weight_ = self._cuda(weights["mlp1.1.weight"][split_start:split_end, :]).t()
+            self.mlp1_1_weight_.copy_(weights["mlp1.1.weight"][split_start:split_end, :].t())
         if "mlp1.1.bias" in weights:
-            self.mlp1_1_bias_ = self._cuda(weights["mlp1.1.bias"][split_start:split_end])
+            self.mlp1_1_bias_.copy_(weights["mlp1.1.bias"][split_start:split_end])
 
         if "mlp1.3.weight" in weights:
-            self.mlp1_3_weight_ = self._cuda(weights["mlp1.3.weight"][:, split_start:split_end]).t()
+            self.mlp1_3_weight_.copy_(weights["mlp1.3.weight"][:, split_start:split_end].t())
         if "mlp1.3.bias" in weights:
-            self.mlp1_3_bias_ = self._cuda(weights["mlp1.3.bias"])
+            self.mlp1_3_bias_.copy_(weights["mlp1.3.bias"])
 
         return
 
diff --git a/lightllm/models/vit/layer_weights/transformer_layer_weight.py b/lightllm/models/vit/layer_weights/transformer_layer_weight.py
index c6024594e3..54ad367867 100644
--- a/lightllm/models/vit/layer_weights/transformer_layer_weight.py
+++ b/lightllm/models/vit/layer_weights/transformer_layer_weight.py
@@ -7,15 +7,16 @@
 from lightllm.common.basemodel.layer_weights.meta_weights import (
     ROWMMWeight,
     COLMMWeight,
-    NoTpNormWeight,
-    TpVitPadNormWeight,
+    RMSNormWeight,
+    LayerNormWeight,
+    TpRMSNormWeight,
 )
 from lightllm.utils.dist_utils import get_current_device_id
 
 
 class ViTTransformerLayerWeight(TransformerLayerWeight):
-    def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+    def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+        super().__init__(layer_num, data_type, network_config, quant_cfg)
         return
 
     def _cuda(self, cpu_tensor):
@@ -29,6 +30,9 @@ def _parse_config(self):
         self.qkv_bias = self.network_config_.get("qkv_bias", True)
         self.layer_norm_eps = self.network_config_.get("layer_norm_eps", 1e-6)
         self.norm_type = self.network_config_.get("norm_type", "layer_norm")
+        self.n_embed = self.network_config_["hidden_size"] + self.padding_hidden_size
+        mlp_ratio = self.network_config_.get("mlp_ratio", 4)
+        self.n_inter = self.network_config_.get("intermediate_size", int(self.n_embed * mlp_ratio))
 
     def _init_weight_names(self):
         self._att_norm_weight_name = f"vision_model.encoder.layers.{self.layer_num_}.norm1.weight"
@@ -81,54 +85,71 @@ def _init_weight(self):
 
     def _init_qkv(self):
         self.qkv_proj = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_embed, self.n_embed, self.n_embed],
             weight_names=[self._q_weight_name, self._k_weight_name, self._v_weight_name],
             data_type=self.data_type_,
             bias_names=[self._q_bias_name, self._k_bias_name, self._v_bias_name],
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="qkv_proj",
+            quant_method=self.get_quant_method("qkv_proj"),
         )
 
     def _init_o(self):
         self.o_proj = COLMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_embed],
             weight_names=self._o_weight_name,
             data_type=self.data_type_,
             bias_names=self._o_bias_name,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="o_proj",
+            quant_method=self.get_quant_method("o_proj"),
         )
 
     def _init_ffn(self):
         self.ffn_1_proj_ = ROWMMWeight(
+            in_dim=self.n_embed,
+            out_dims=[self.n_inter],
             weight_names=self.fc1_weight_name_,
             data_type=self.data_type_,
             bias_names=self.fc1_bias_name_,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="ffn_1_proj",
+            quant_method=self.get_quant_method("ffn_1_proj"),
         )
 
         self.ffn_2_proj_ = COLMMWeight(
+            in_dim=self.n_inter,
+            out_dims=[self.n_embed],
             weight_names=self.fc2_weight_name_,
             data_type=self.data_type_,
             bias_names=self.fc2_bias_name_,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="ffn_2_proj",
+            quant_method=self.get_quant_method("ffn_2_proj"),
         )
 
     def _init_norm(self):
-        self.att_norm_weight_ = NoTpNormWeight(
-            self._att_norm_weight_name, self.data_type_, bias_name=self._att_norm_bias_name
+        norm_weight_cls = RMSNormWeight if self.norm_type == "rms_norm" else LayerNormWeight
+        hidden_size = self.network_config_["hidden_size"]
+        self.att_norm_weight_ = norm_weight_cls(
+            dim=hidden_size,
+            weight_name=self._att_norm_weight_name,
+            data_type=self.data_type_,
+            bias_name=self._att_norm_bias_name,
         )
-        self.ffn_norm_weight_ = NoTpNormWeight(
-            self._ffn_norm_weight_name, self.data_type_, bias_name=self._ffn_norm_bias_name
+        self.ffn_norm_weight_ = norm_weight_cls(
+            dim=hidden_size,
+            weight_name=self._ffn_norm_weight_name,
+            data_type=self.data_type_,
+            bias_name=self._ffn_norm_bias_name,
         )
         if self.qk_norm:
-            head_num = self.network_config_["num_attention_heads"]
-            self.q_norm_weight_ = TpVitPadNormWeight(self._q_norm_weight_name, self.data_type_, head_num=head_num)
-            self.k_norm_weight_ = TpVitPadNormWeight(self._k_norm_weight_name, self.data_type_, head_num=head_num)
+            self.q_norm_weight_ = TpRMSNormWeight(
+                dim=hidden_size,
+                weight_name=self._q_norm_weight_name,
+                data_type=self.data_type_,
+                bias_name=None,
+            )
+            self.k_norm_weight_ = TpRMSNormWeight(
+                dim=hidden_size,
+                weight_name=self._k_norm_weight_name,
+                data_type=self.data_type_,
+                bias_name=None,
+            )
 
     def load_hf_weights(self, weights):
         if f"vision_model.encoder.layers.{self.layer_num_}.attn.qkv.weight" in weights:
diff --git a/lightllm/models/vit/model.py b/lightllm/models/vit/model.py
index 01bb69bdfe..9c2bc42426 100644
--- a/lightllm/models/vit/model.py
+++ b/lightllm/models/vit/model.py
@@ -40,7 +40,6 @@ def __init__(self, kvargs):
         self.tp_world_size_ = get_dp_world_size()
         self.weight_dir_ = kvargs["weight_dir"]
         self.load_way = kvargs.get("load_way", "HF")
-        self.mode = [m.replace("int4weight", "w4a16").replace("int8weight", "w8a16") for m in kvargs.get("mode", [])]
         self.weight_dict = kvargs.get("weight_dict", None)
         self.data_type = kvargs.get("data_type", "float16")
         self.quant_type = kvargs.get("quant_type", None)
@@ -112,15 +111,12 @@ def _padding_hidden_size(self):
         return
 
     def _init_weights(self):
-        self.pre_post_weight = self.pre_and_post_weight_class(
-            self.data_type, network_config=self.config, mode=self.mode
-        )
+        self.pre_post_weight = self.pre_and_post_weight_class(self.data_type, network_config=self.config)
         self.trans_layers_weight = [
             self.transformer_weight_class(
                 i,
                 self.data_type,
                 network_config=self.config,
-                mode=self.mode,
                 quant_cfg=self.quant_cfg,
             )
             for i in range(self.config["num_hidden_layers"])
@@ -141,10 +137,10 @@ def _init_quant(self):
         logger.info(f"Initial quantization. " f"The default quantization method is {self.quant_cfg.quant_type}")
 
     def _init_infer_layer(self):
-        self.pre_infer = self.pre_layer_infer_class(network_config=self.config, mode=self.mode)
-        self.post_infer = self.post_layer_infer_class(network_config=self.config, mode=self.mode)
+        self.pre_infer = self.pre_layer_infer_class(network_config=self.config)
+        self.post_infer = self.post_layer_infer_class(network_config=self.config)
         self.layers_infer = [
-            self.transformer_layer_infer_class(i, network_config=self.config, mode=self.mode)
+            self.transformer_layer_infer_class(i, network_config=self.config)
             for i in range(self.config["num_hidden_layers"])
         ]
         return
@@ -185,7 +181,7 @@ def encode(self, images: List[ImageItem]):
             else:
                 raise Exception("Unsupport input types: {} for {}".format(type(img), img))
 
-            cur_num = img_tensors[-1].shape[0]
+            cur_num = img.token_num
             valid_ids.append([valid_id, valid_id + cur_num])
             valid_id += cur_num
 
@@ -195,7 +191,7 @@ def encode(self, images: List[ImageItem]):
         imgs = torch.cat(img_tensors, dim=0)
         pixel_values = imgs.cuda().to(dtype=self.data_type)
         all_img_embeds = self.forward(pixel_values)
-        return all_img_embeds, uuids, valid_ids
+        return all_img_embeds.view(-1, all_img_embeds.shape[-1]), uuids, valid_ids
 
     def cuda(self):
         return self
diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py
index d193bab41c..82ff934554 100644
--- a/lightllm/server/api_cli.py
+++ b/lightllm/server/api_cli.py
@@ -219,29 +219,6 @@ def make_argument_parser() -> argparse.ArgumentParser:
         the --nccl_host must equal to the config_server_host, and the --nccl_port must be unique for a config_server,
         dont use same nccl_port for different inference node, it will be critical error""",
     )
-
-    parser.add_argument(
-        "--mode",
-        type=str,
-        default=[],
-        nargs="+",
-        help="""Model mode: [triton_int8kv | ppl_int8kv | ppl_int8kv_flashdecoding | ppl_int8kv_flashdecoding_diverse
-                        | ppl_fp16 | triton_flashdecoding
-                        | triton_gqa_attention | triton_gqa_flashdecoding | triton_fp8kv | offline_calibration_fp8kv
-                        | export_fp8kv_calibration
-                        triton_flashdecoding mode is for long context, current support llama llama2 qwen;
-                        triton_gqa_attention and triton_gqa_flashdecoding is fast kernel for model which use GQA;
-                        triton_int8kv mode use int8 to store kv cache, can increase token capacity, use triton kernel;
-                        triton_fp8kv mode use float8 to store kv cache, currently only for deepseek2;
-                        offline_calibration_fp8kv mode use float8 to store kv cache, need fa3 or flashinfer backend,
-                        currently only for llama and qwen model;
-                        export_fp8kv_calibration record and export kv cache quant calibration results to a json file.
-                        It can be used for llama and qwen model.
-                        Calibration need to disable cudagraph and use fa3 or flashinfer backend.
-                        ppl_int8kv mode use int8 to store kv cache, and use ppl fast kernel;
-                        ppl_fp16 mode use ppl fast fp16 decode attention kernel;
-                        you need to read source code to make sure the supported detail mode for all models""",
-    )
     parser.add_argument(
         "--trust_remote_code",
         action="store_true",
@@ -337,21 +314,40 @@ def make_argument_parser() -> argparse.ArgumentParser:
         only deepseekv3 model supported now.""",
     )
     parser.add_argument(
-        "--enable_flashinfer_prefill",
-        action="store_true",
-        help="""inference backend will use the attention kernel of flashinfer for prefill,
-        only deepseekv3 model supported now.""",
+        "--llm_prefill_att_backend",
+        type=str,
+        nargs="+",
+        choices=["None", "triton", "fa3", "flashinfer"],
+        default=["fa3"],
+        help="""prefill attention kernel used in llm.
+                None: automatically select backend based on current GPU device,
+                not supported yet, will support in future""",
     )
     parser.add_argument(
-        "--enable_flashinfer_decode",
-        action="store_true",
-        help="""inference backend will use the attention kernel of flashinfer for decode,
-        only deepseekv3 model supported now.""",
+        "--llm_decode_att_backend",
+        type=str,
+        nargs="+",
+        choices=["None", "triton", "fa3", "flashinfer"],
+        default=["fa3"],
+        help="""decode attention kernel used in llm.
+                None: automatically select backend based on current GPU device,
+                not supported yet, will support in future""",
     )
     parser.add_argument(
-        "--enable_fa3",
-        action="store_true",
-        help="""inference backend will use the fa3 attention kernel for prefill and decode""",
+        "--llm_kv_type",
+        type=str,
+        choices=["None", "int8kv", "int4kv"],
+        default="None",
+        help="""kv type used in llm, None for dtype that llm used in config.json.
+                fp8kv: not fully supported yet, will support in future""",
+    )
+    parser.add_argument(
+        "--llm_kv_quant_group_size",
+        type=int,
+        default=8,
+        help="""kv quant group size used in llm kv, when llm_kv_type is quanted type,such as int8kv,
+        this params will be effective.
+        """,
     )
     parser.add_argument(
         "--cache_capacity", type=int, default=200, help="cache server capacity for multimodal resources"
@@ -469,10 +465,8 @@ def make_argument_parser() -> argparse.ArgumentParser:
         "--quant_type",
         type=str,
         default="none",
-        help="""Quantization method: ppl-w4a16-128 | flashllm-w6a16
-                        | ao-int4wo-[32,64,128,256] | ao-int8wo | ao-fp8w8a16 | ao-fp6w6a16
-                        | vllm-w8a8 | vllm-fp8w8a8 | vllm-fp8w8a8-b128
-                        | triton-fp8w8a8-block128""",
+        help="""Quantization method: vllm-w8a8 | vllm-fp8w8a8 | vllm-fp8w8a8-b128
+                        | deepgemm-fp8w8a8-b128 | triton-fp8w8a8-block128 | awq | awq_marlin""",
     )
     parser.add_argument(
         "--quant_cfg",
@@ -485,9 +479,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
         "--vit_quant_type",
         type=str,
         default="none",
-        help="""Quantization method: ppl-w4a16-128 | flashllm-w6a16
-                        | ao-int4wo-[32,64,128,256] | ao-int8wo | ao-fp8w8a16 | ao-fp6w6a16
-                        | vllm-w8a8 | vllm-fp8w8a8""",
+        help="""Quantization method for ViT: vllm-w8a8 | vllm-fp8w8a8""",
     )
     parser.add_argument(
         "--vit_quant_cfg",
@@ -524,6 +516,11 @@ def make_argument_parser() -> argparse.ArgumentParser:
             " Therefore, it is recommended to set this parameter according to actual needs."
         ),
     )
+    parser.add_argument(
+        "--enable_ep_moe",
+        action="store_true",
+        help="""Whether to enable ep moe for deepseekv3 model.""",
+    )
     parser.add_argument(
         "--ep_redundancy_expert_config_path",
         type=str,
@@ -538,7 +535,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--enable_fused_shared_experts",
         action="store_true",
-        help="""Whether to enable fused shared experts for deepseekv3 model. only work when MOE_MODE=TP """,
+        help="""Whether to enable fused shared experts for deepseekv3 model. only work when tensor parallelism""",
     )
     parser.add_argument(
         "--mtp_mode",
@@ -612,4 +609,24 @@ def make_argument_parser() -> argparse.ArgumentParser:
         default=False,
         help="""Enable prefix prompt cache fetch for data parallel inference, disabled by default.""",
     )
+    parser.add_argument(
+        "--hardware_platform",
+        type=str,
+        default="cuda",
+        choices=["cuda", "musa"],
+        help="""Hardware platform: cuda | musa""",
+    )
+    parser.add_argument(
+        "--enable_torch_fallback",
+        action="store_true",
+        help="""Whether to enable torch naive implementation for the op.
+        If the op is not implemented for the platform, it will use torch naive implementation.""",
+    )
+    parser.add_argument(
+        "--enable_triton_fallback",
+        action="store_true",
+        help="""Whether to enable triton implementation for the op.
+        If the op is not implemented for the platform and the hardware support triton,
+        it will use triton implementation.""",
+    )
     return parser
diff --git a/lightllm/server/api_models.py b/lightllm/server/api_models.py
index 7b9cdd5012..f30ecc55fe 100644
--- a/lightllm/server/api_models.py
+++ b/lightllm/server/api_models.py
@@ -24,9 +24,10 @@ class Message(BaseModel):
 class Function(BaseModel):
     """Function descriptions."""
 
-    description: Optional[str] = Field(default=None, examples=[None])
     name: Optional[str] = None
-    parameters: Optional[object] = None
+    description: Optional[str] = Field(default=None, examples=[None])
+    parameters: Optional[dict] = None
+    response: Optional[dict] = None
 
 
 class Tool(BaseModel):
diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py
index 6a8c232dc5..d91bb1d947 100644
--- a/lightllm/server/api_openai.py
+++ b/lightllm/server/api_openai.py
@@ -81,7 +81,7 @@ def _process_tool_call_id(
         # SGLang sets call_item.tool_index to the *local* position inside that message.
         # Therefore, the index must be corrected by using
         # `history_tool_calls_cnt + call_item.tool_index` to ensure globally unique and properly ordered.
-        tool_call_id = f"functions.{call_item.name}:{history_tool_calls_cnt+call_item.tool_index}"
+        tool_call_id = f"functions.{call_item.name}:{history_tool_calls_cnt + call_item.tool_index}"
         logger.debug(
             f"Process tool call idx, parser: {tool_call_parser}, \
             tool_call_id: {tool_call_id}, \
diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py
index 4ead3cbbf7..3ae3789f47 100644
--- a/lightllm/server/api_start.py
+++ b/lightllm/server/api_start.py
@@ -122,21 +122,6 @@ def normal_or_p_d_start(args):
     if args.return_all_prompt_logprobs:
         assert args.disable_dynamic_prompt_cache is True, "need add --disable_dynamic_prompt_cache"
         assert args.disable_chunked_prefill is True, "need add --disable_chunked_prefill"
-    if "offline_calibration_fp8kv" in args.mode:
-        assert args.enable_fa3 is True or (
-            args.enable_flashinfer_prefill is True and args.enable_flashinfer_decode is True
-        ), (
-            "offline_calibration_fp8kv mode need enable fa3 or flashinfer, add --enable_fa3 or "
-            "--enable_flashinfer_prefill and --enable_flashinfer_decode"
-        )
-    if "export_fp8kv_calibration" in args.mode:
-        assert args.enable_fa3 is True or (
-            args.enable_flashinfer_prefill is True and args.enable_flashinfer_decode is True
-        ), (
-            "export_fp8kv_calibration mode need enable fa3 or flashinfer, add --enable_fa3 or "
-            "--enable_flashinfer_prefill and --enable_flashinfer_decode"
-        )
-        assert args.disable_cudagraph is True, "export_fp8kv_calibration mode need disable cudagraph"
 
     # 部分模式还不能支持与高级动态调度算法协同，to do.
     if args.diverse_mode:
diff --git a/lightllm/server/core/objs/sampling_params.py b/lightllm/server/core/objs/sampling_params.py
index f073319d79..d955aa6a87 100644
--- a/lightllm/server/core/objs/sampling_params.py
+++ b/lightllm/server/core/objs/sampling_params.py
@@ -7,7 +7,7 @@
 
 _SAMPLING_EPS = 1e-5
 DEFAULT_INPUT_PENALTY = os.getenv("INPUT_PENALTY", "False").upper() in ["ON", "TRUE", "1"]
-SKIP_SPECIAL_TOKENS = os.getenv("SKIP_SPECIAL_TOKENS", "True").upper() in ["ON", "TRUE", "1"]
+SKIP_SPECIAL_TOKENS = os.getenv("SKIP_SPECIAL_TOKENS", "False").upper() in ["ON", "TRUE", "1"]
 
 # 从环境变量获取最大长度限制
 STOP_SEQUENCE_MAX_LENGTH = int(os.getenv("LIGHTLLM_STOP_SEQUENCE_MAX_LENGTH", 256))
diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py
index 5ebadaf165..7b440718f5 100644
--- a/lightllm/server/core/objs/start_args_type.py
+++ b/lightllm/server/core/objs/start_args_type.py
@@ -63,7 +63,6 @@ class StartArgs:
     nccl_host: str = field(default="127.0.0.1")
     nccl_port: int = field(default=28765)
     use_config_server_to_init_nccl: bool = field(default=False)
-    mode: List[str] = field(default_factory=list)
     trust_remote_code: bool = field(default=False)
     disable_log_stats: bool = field(default=False)
     log_stats_interval: int = field(default=10)
@@ -116,8 +115,14 @@ class StartArgs:
     quant_cfg: Optional[str] = field(default=None)
     vit_quant_type: Optional[str] = field(default=None)
     vit_quant_cfg: Optional[str] = field(default=None)
-    enable_flashinfer_prefill: bool = field(default=False)
-    enable_flashinfer_decode: bool = field(default=False)
+    llm_prefill_att_backend: List[str] = field(
+        default=("None",), metadata={"choices": ["None", "triton", "fa3", "flashinfer"]}
+    )
+    llm_decode_att_backend: List[str] = field(
+        default=("None",), metadata={"choices": ["None", "triton", "fa3", "flashinfer"]}
+    )
+    llm_kv_type: str = field(default="None", metadata={"choices": ["None", "int8kv", "int4kv", "fp8kv"]})
+    llm_kv_quant_group_size: int = field(default=8)
     sampling_backend: str = field(default="triton", metadata={"choices": ["triton", "sglang_kernel"]})
     penalty_counter_mode: str = field(
         default="gpu_counter", metadata={"choices": ["cpu_counter", "pin_mem_counter", "gpu_counter"]}
@@ -153,6 +158,3 @@ class StartArgs:
     # multi_modal
     enable_multimodal: bool = field(default=False)
     enable_multimodal_audio: bool = field(default=False)
-
-    # kernel setting
-    enable_fa3: bool = field(default=False)
diff --git a/lightllm/server/router/manager.py b/lightllm/server/router/manager.py
index 89c46d9ed9..ac5c1abee3 100644
--- a/lightllm/server/router/manager.py
+++ b/lightllm/server/router/manager.py
@@ -58,7 +58,6 @@ def __init__(self, args: StartArgs):
         # 判断是否是保守调度，保守调度不会发生暂停 req 的情况，但是有些场景可能影响吞吐
         self.is_safe_schedule = args.router_token_ratio == 0.0
         self.load_way = args.load_way
-        self.mode = args.mode
         self.max_total_token_num = args.max_total_token_num
         self.shm_req_manager = ShmReqManager()
         # 用共享内存进行共享，router 模块读取进行精确的调度估计
@@ -155,7 +154,6 @@ async def wait_to_model_ready(self):
             "weight_dir": self.model_weightdir,
             "load_way": self.load_way,
             "max_total_token_num": self.max_total_token_num,
-            "mode": self.mode,
             "max_req_num": self.args.running_max_req_size + 8,
             "max_seq_length": self.args.max_req_total_len + 8,  # 留一点余量
             "nccl_host": self.args.nccl_host,
diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py
index 92653bc0cd..805c9b8e50 100644
--- a/lightllm/server/router/model_infer/mode_backend/base_backend.py
+++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py
@@ -88,7 +88,6 @@ def init_model(self, kvargs):
         # dp_size_in_node 计算兼容多机纯tp的运行模式，这时候 1 // 2 == 0, 需要兼容
         self.dp_size_in_node = max(1, self.dp_size // self.nnodes)
         self.load_way = kvargs["load_way"]
-        self.mode = kvargs["mode"]
         self.disable_chunked_prefill = self.args.disable_chunked_prefill
         self.chunked_prefill_size = self.args.chunked_prefill_size
         self.return_all_prompt_logprobs = self.args.return_all_prompt_logprobs
@@ -148,7 +147,6 @@ def init_model(self, kvargs):
             "weight_dir": self.weight_dir,
             "max_total_token_num": max_total_token_num,
             "load_way": self.load_way,
-            "mode": self.mode,
             "max_req_num": kvargs.get("max_req_num", 1000),
             "max_seq_length": kvargs.get("max_seq_length", 1024 * 5),
             "is_token_healing": kvargs.get("is_token_healing", False),
@@ -302,7 +300,6 @@ def init_mtp_draft_model(self, main_kvargs: dict):
                 "weight_dir": self.args.mtp_draft_model_dir[i],
                 "max_total_token_num": self.model.mem_manager.size,
                 "load_way": main_kvargs["load_way"],
-                "mode": main_kvargs["mode"],
                 "max_req_num": main_kvargs.get("max_req_num", 1000),
                 "max_seq_length": main_kvargs.get("max_seq_length", 1024 * 5),
                 "is_token_healing": False,
diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py
index f3450261b7..a8a5224ebc 100644
--- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py
+++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py
@@ -399,7 +399,7 @@ def _draft_decode_eagle(
             draft_model_output: ModelOutput = self.draft_models[draft_model_idx].forward(draft_model_input)
             draft_next_token_ids = self._gen_argmax_token_ids(draft_model_output)
             draft_model_input.b_seq_len += 1
-            draft_model_input.max_len_in_batch += 1
+            draft_model_input.max_kv_seq_len += 1
             eagle_mem_indexes_i = eagle_mem_indexes[_step * num_reqs : (_step + 1) * num_reqs]
             draft_model_input.mem_indexes = torch.cat(
                 [draft_model_input.mem_indexes.view(-1, self.mtp_step + 1)[:, 1:], eagle_mem_indexes_i.view(-1, 1)],
diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py
index df10a6d4e6..bb0e848e76 100644
--- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py
+++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py
@@ -591,7 +591,7 @@ def _draft_decode_eagle(
             draft_model_output: ModelOutput = self.draft_models[draft_model_idx].forward(draft_model_input)
             # update the meta info of the inference
             draft_model_input.b_seq_len += 1
-            draft_model_input.max_len_in_batch += 1
+            draft_model_input.max_kv_seq_len += 1
             eagle_mem_indexes_i = eagle_mem_indexes[_step * real_req_num : (_step + 1) * real_req_num]
             eagle_mem_indexes_i = F.pad(
                 input=eagle_mem_indexes_i,
@@ -955,7 +955,7 @@ def _draft_decode_eagle_overlap(
             )
 
             draft_model_input0.b_seq_len += 1
-            draft_model_input0.max_len_in_batch += 1
+            draft_model_input0.max_kv_seq_len += 1
             eagle_mem_indexes_i = eagle_mem_indexes0[_step * real_req_num0 : (_step + 1) * real_req_num0]
             eagle_mem_indexes_i = F.pad(
                 input=eagle_mem_indexes_i,
@@ -969,7 +969,7 @@ def _draft_decode_eagle_overlap(
             ).view(-1)
 
             draft_model_input1.b_seq_len += 1
-            draft_model_input1.max_len_in_batch += 1
+            draft_model_input1.max_kv_seq_len += 1
             eagle_mem_indexes_i = eagle_mem_indexes1[_step * real_req_num1 : (_step + 1) * real_req_num1]
             eagle_mem_indexes_i = F.pad(
                 input=eagle_mem_indexes_i,
diff --git a/lightllm/server/router/model_infer/mode_backend/generic_padded_pre_process.py b/lightllm/server/router/model_infer/mode_backend/generic_padded_pre_process.py
index 6465995c45..03ac4cfb05 100644
--- a/lightllm/server/router/model_infer/mode_backend/generic_padded_pre_process.py
+++ b/lightllm/server/router/model_infer/mode_backend/generic_padded_pre_process.py
@@ -73,7 +73,6 @@ def padded_prepare_prefill_inputs(
     max_kv_seq_len = max(b_seq_len)
     max_cache_len = max(b_ready_cache_len)
     max_q_seq_len = max(b_q_seq_len)
-    max_len_in_batch = max(b_q_seq_len)
 
     input_ids = np.concatenate(input_ids, dtype=np.int64)
     input_ids = torch.tensor(input_ids, dtype=torch.int64, device="cpu")
@@ -102,7 +101,6 @@ def padded_prepare_prefill_inputs(
     model_input = ModelInput(
         batch_size=b_seq_len.shape[0],
         total_token_num=total_token_num,
-        max_len_in_batch=max_len_in_batch,
         max_q_seq_len=max_q_seq_len,
         max_kv_seq_len=max_kv_seq_len,
         max_cache_len=max_cache_len,
@@ -150,6 +148,7 @@ def padded_prepare_decode_inputs(
         seq_len = req.get_cur_total_len()
         assert req.cur_kv_len == seq_len - 1
         b_seq_len.append(seq_len)
+        b_q_seq_len.append(1)
         total_token_num += seq_len
         b_mtp_index.append(0)
         batch_multimodal_params.append(req.multimodal_params)
@@ -160,32 +159,30 @@ def padded_prepare_decode_inputs(
             total_token_num += seq_len
             b_req_idx.append(req.req_idx)
             b_seq_len.append(seq_len)
+            b_q_seq_len.append(1)
             b_mtp_index.append(step + 1)
             batch_multimodal_params.append(req.multimodal_params)
 
-        b_q_seq_len.append(req.mtp_step + 1)
-
     # padding fake req for decode
     for _ in range(padded_req_num):
         seq_len = 2
         total_token_num += seq_len
         b_req_idx.append(g_infer_context.req_manager.HOLD_REQUEST_ID)
         b_seq_len.append(seq_len)
+        b_q_seq_len.append(1)
         b_mtp_index.append(0)
         batch_multimodal_params.append({"images": [], "audios": []})
         for step in range(args_mtp_step):
             seq_len += 1
             total_token_num += seq_len
             b_seq_len.append(seq_len)
+            b_q_seq_len.append(1)
             b_req_idx.append(g_infer_context.req_manager.HOLD_REQUEST_ID)
             b_mtp_index.append(step + 1)
             batch_multimodal_params.append({"images": [], "audios": []})
 
-        b_q_seq_len.append(1 + args_mtp_step)
-
     max_kv_seq_len = max(b_seq_len)
     max_q_seq_len = max(b_q_seq_len)
-    max_len_in_batch = max(b_seq_len)
 
     b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu")
     b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu")
@@ -210,7 +207,6 @@ def padded_prepare_decode_inputs(
     model_input = ModelInput(
         batch_size=b_seq_len.shape[0],
         total_token_num=total_token_num,
-        max_len_in_batch=max_len_in_batch,
         max_q_seq_len=max_q_seq_len,
         max_kv_seq_len=max_kv_seq_len,
         input_ids=None,
diff --git a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py
index bdb36054b4..4eb8c7e1e6 100644
--- a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py
+++ b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py
@@ -50,7 +50,6 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) ->
 
     max_kv_seq_len = max(b_seq_len)
     max_cache_len = max(b_ready_cache_len)
-    max_len_in_batch = max(b_q_seq_len)
     max_q_seq_len = max(b_q_seq_len)
 
     input_ids = np.concatenate(input_ids, dtype=np.int64)
@@ -72,7 +71,6 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) ->
     model_input = ModelInput(
         batch_size=b_seq_len.shape[0],
         total_token_num=total_token_num,
-        max_len_in_batch=max_len_in_batch,
         max_q_seq_len=max_q_seq_len,
         max_kv_seq_len=max_kv_seq_len,
         max_cache_len=max_cache_len,
@@ -95,7 +93,6 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) ->
 def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[InferReq]]:
     run_reqs: List[InferReq] = []
     total_token_num = 0
-    max_len_in_batch = 0
     b_req_idx = []
     b_mtp_index = []
     b_seq_len = []
@@ -107,8 +104,8 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
         seq_len = req.get_cur_total_len()
         assert req.cur_kv_len == seq_len - 1, f"{req.cur_kv_len} {seq_len}"
         b_seq_len.append(seq_len)
+        b_q_seq_len.append(1)
         total_token_num += seq_len
-        max_len_in_batch = max(max_len_in_batch, seq_len)
         b_mtp_index.append(0)
         multimodal_params.append(req.multimodal_params)
         # process the draft tokens.
@@ -118,10 +115,9 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
             seq_len += 1
             b_seq_len.append(seq_len)
             total_token_num += seq_len
-            max_len_in_batch = max(max_len_in_batch, seq_len)
             b_mtp_index.append(step + 1)
             multimodal_params.append(req.multimodal_params)
-        b_q_seq_len.append(req.mtp_step + 1)
+            b_q_seq_len.append(1)
 
     max_kv_seq_len = max(b_seq_len)
     max_q_seq_len = max(b_q_seq_len)
@@ -146,7 +142,6 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
     model_input = ModelInput(
         batch_size=b_seq_len.shape[0],
         total_token_num=total_token_num,
-        max_len_in_batch=max_len_in_batch,
         max_q_seq_len=max_q_seq_len,
         max_kv_seq_len=max_kv_seq_len,
         input_ids=None,
diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py
index d4ba902999..b9d13f512b 100644
--- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py
+++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py
@@ -26,11 +26,16 @@ def __init__(self, backend):
         self.init_sync_group = create_new_group_for_current_dp("nccl")
         dist.barrier(group=self.init_sync_group)
 
+        self.page_index_buffer = torch.empty((1024 * 1024 * 4,), dtype=torch.int32, device="cuda")
+        self.page_ready_buffer = torch.empty((1024 * 1024 * 4,), dtype=torch.bool, device="cuda")
+
         self.cpu_cache_handle_queue: Deque[TransTask] = deque()
         self.cpu_cache_client = CpuKvCacheClient(only_create_meta_data=False, init_shm_data=False)
 
         # 一些算子模式需要同步计算和 cpu cache 的 load 和 offload 操作
-        self.need_sync_compute_stream: bool = True
+        self.need_sync_compute_stream: bool = (
+            "fa3" in self.args.llm_decode_att_backend or "fa3" in self.args.llm_prefill_att_backend
+        )
 
     def wait(self):
         """
@@ -89,14 +94,18 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]):
                         cpu_kv_cache_scale = None
                         gpu_kv_cache_scale = None
 
+                    mem_indexes_cuda = mem_indexes.cuda(non_blocking=True)
+                    page_indexes_cuda = torch.tensor(need_pages, dtype=torch.int32, device="cpu").cuda(
+                        non_blocking=True
+                    )
                     # 将 cpu page 的内容拷贝到 gpu 页面中
                     load_cpu_kv_to_gpu(
-                        gpu_mem_indexes=mem_indexes.cuda(non_blocking=True),
+                        gpu_mem_indexes=mem_indexes_cuda,
                         gpu_kv_cache=mem_manager.kv_buffer,
                         gpu_kv_cache_scale=gpu_kv_cache_scale,
                         cpu_kv_cache=cpu_kv_cache,
                         cpu_kv_cache_scale=cpu_kv_cache_scale,
-                        page_indexes=torch.tensor(need_pages, dtype=torch.int32, device="cpu").cuda(non_blocking=True),
+                        page_indexes=page_indexes_cuda,
                         tp_index=self.backend.rank_in_dp,
                         tp_world_size=self.backend.dp_world_size,
                         grid_num=grid_num,
@@ -221,6 +230,12 @@ def _start_kv_cache_offload_task(
 
             page_indexes = torch.tensor(page_list, dtype=torch.int32, device="cpu", pin_memory=True)
             page_readies = torch.tensor(ready_list, dtype=torch.bool, device="cpu", pin_memory=True)
+            assert len(page_indexes) <= self.page_index_buffer.shape[0]
+            cuda_page_indexes = self.page_index_buffer[: len(page_indexes)]
+            cuda_page_readies = self.page_ready_buffer[: len(page_readies)]
+            cuda_page_indexes.copy_(page_indexes, non_blocking=True)
+            cuda_page_readies.copy_(page_readies, non_blocking=True)
+
             move_token_num = item_size * self.args.cpu_cache_token_page_size
             assert req.cur_kv_len >= item_size * self.args.cpu_cache_token_page_size
             token_indexes = self.backend.model.req_manager.req_to_token_indexs[req.req_idx, 0:move_token_num]
@@ -248,8 +263,8 @@ def _start_kv_cache_offload_task(
                 gpu_kv_cache_scale=gpu_kv_cache_scale,
                 cpu_kv_cache=cpu_kv_cache,
                 cpu_kv_cache_scale=cpu_kv_cache_scale,
-                page_indexes=page_indexes,
-                page_readies=page_readies,
+                page_indexes=cuda_page_indexes,
+                page_readies=cuda_page_readies,
                 tp_index=self.backend.rank_in_dp,
                 tp_world_size=self.backend.dp_world_size,
                 grid_num=grid_num,
diff --git a/lightllm/server/router/model_infer/mode_backend/redundancy_expert_manager.py b/lightllm/server/router/model_infer/mode_backend/redundancy_expert_manager.py
index 811d39a729..596eca4f24 100644
--- a/lightllm/server/router/model_infer/mode_backend/redundancy_expert_manager.py
+++ b/lightllm/server/router/model_infer/mode_backend/redundancy_expert_manager.py
@@ -8,10 +8,10 @@
 import json
 from typing import List
 from lightllm.common.basemodel.basemodel import TpPartBaseModel
-from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe_weight_ep_redundancy import (
+from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_redundancy import (
     FusedMoeWeightEPAutoRedundancy,
 )
-from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe_weight_ep import FusedMoeWeightEP
+from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight
 from lightllm.utils.envs_utils import get_env_start_args, get_redundancy_expert_update_interval
 from lightllm.utils.envs_utils import get_redundancy_expert_update_max_load_count
 from lightllm.utils.envs_utils import get_redundancy_expert_num
@@ -28,7 +28,7 @@ def __init__(self, model: TpPartBaseModel):
         self.model = model
         self.ep_fused_moeweights: List[FusedMoeWeightEPAutoRedundancy] = []
         for layer in self.model.trans_layers_weight:
-            ep_weights = self._find_members_of_class(layer, FusedMoeWeightEP)
+            ep_weights = self._find_members_of_class(layer, FusedMoeWeight)
             assert len(ep_weights) <= 1
             self.ep_fused_moeweights.extend([FusedMoeWeightEPAutoRedundancy(e) for e in ep_weights])
 
diff --git a/lightllm/utils/device_utils.py b/lightllm/utils/device_utils.py
index cd48a355bd..a1ed6ed950 100644
--- a/lightllm/utils/device_utils.py
+++ b/lightllm/utils/device_utils.py
@@ -3,6 +3,8 @@
 import torch
 import shutil
 import subprocess
+from enum import Enum
+from typing import Optional
 from functools import lru_cache
 from lightllm.utils.log_utils import init_logger
 
@@ -81,11 +83,14 @@ def calcu_kernel_best_vsm_count(kernel, num_warps):
     return num_sm
 
 
+@lru_cache(maxsize=1)
+def is_musa():
+    return hasattr(torch.version, "musa") and torch.version.musa is not None
+
+
 @lru_cache(maxsize=None)
 def get_current_device_name():
-    import torch
-
-    if torch.cuda.is_available():
+    if torch.cuda.is_available() or is_musa():
         device = torch.cuda.current_device()
         gpu_name = torch.cuda.get_device_name(device)
         # 4090 trans to 4090 D
@@ -103,8 +108,6 @@ def init_p2p(device_index):
     """
     torch 调用跨卡的to操作后，triton编译的算子便能自动操作跨卡tensor。
     """
-    import torch
-
     num_gpus = torch.cuda.device_count()
     tensor = torch.zeros((1,))
     tensor = tensor.to(f"cuda:{device_index}")
@@ -127,8 +130,26 @@ def has_nvlink():
         result = result.decode("utf-8")
         # Check if the output contains 'NVLink'
         return any(f"NV{i}" in result for i in range(1, 8))
+    except FileNotFoundError:
+        # nvidia-smi is not installed, assume no NVLink
+        return False
     except subprocess.CalledProcessError:
-        # If there's an error (e.g., nvidia-smi is not installed or another issue), assume no NVLink
+        # If there's an error while executing nvidia-smi, assume no NVLink
+        return False
+
+
+def has_mtlink():
+    try:
+        # Call mthreads-gmi to get the topology matrix
+        result = subprocess.check_output(["mthreads-gmi", "topo", "--matrix"])
+        result = result.decode("utf-8")
+        # Check if the output contains 'MTLink'
+        return any(f"MT{i}" in result for i in range(1, 8))
+    except FileNotFoundError:
+        # mthreads-gmi is not installed, assume no MTLink
+        return False
+    except subprocess.CalledProcessError:
+        # If there's an error while executing mthreads-gmi, assume no MTLink
         return False
 
 
@@ -265,3 +286,42 @@ def is_5090_gpu() -> bool:
             return False
     except:
         return False
+
+
+class Platform(Enum):
+    """hardware platform enum"""
+
+    CUDA = "cuda"
+    ASCEND = "ascend"  # ascend
+    CAMBRICON = "cambricon"  # cambricon
+    MUSA = "musa"  # musa
+    ROCM = "rocm"  # rocm
+    CPU = "cpu"  # cpu
+
+
+# 目前仅支持cuda 和 musa
+def get_platform(platform_name: Optional[str] = None) -> Platform:
+    """
+    get hardware platform.
+
+    Args:
+        platform_name: platform name (cuda, ascend, cambricon, musa, rocm, cpu)
+
+    Returns:
+        Platform: platform enum value
+    """
+    assert platform_name in ["cuda", "musa"], f"Only support cuda and musa now, but got {platform_name}"
+    platform_name = platform_name.lower()
+    platform_map = {
+        "cuda": Platform.CUDA,
+        "ascend": Platform.ASCEND,
+        "cambricon": Platform.CAMBRICON,
+        "musa": Platform.MUSA,
+        "rocm": Platform.ROCM,
+        "cpu": Platform.CPU,
+    }
+
+    platform = platform_map.get(platform_name)
+    if platform is None:
+        raise ValueError(f"Unknown platform name: {platform_name}")
+    return platform
diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py
index b5822a342c..0a70f1dfa6 100644
--- a/lightllm/utils/envs_utils.py
+++ b/lightllm/utils/envs_utils.py
@@ -26,10 +26,16 @@ def get_unique_server_name():
 def set_cuda_arch(args):
     if not torch.cuda.is_available():
         return
-    if args.enable_flashinfer_prefill or args.enable_flashinfer_decode:
+
+    from lightllm.server.core.objs.start_args_type import StartArgs
+
+    args: StartArgs = args
+
+    if "flashinfer" in args.llm_prefill_att_backend or "flashinfer" in args.llm_decode_att_backend:
         capability = torch.cuda.get_device_capability()
         arch = f"{capability[0]}.{capability[1]}"
         os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}"
+    return
 
 
 def set_env_start_args(args):
@@ -209,7 +215,7 @@ def get_diverse_max_batch_shared_group_size() -> int:
 
 @lru_cache(maxsize=None)
 def enable_diverse_mode_gqa_decode_fast_kernel() -> bool:
-    return get_env_start_args().diverse_mode and "ppl_int8kv_flashdecoding_diverse" in get_env_start_args().mode
+    return get_env_start_args().diverse_mode and "int8kv" == get_env_start_args().llm_kv_type
 
 
 @lru_cache(maxsize=None)
diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py
index 4875b4eee6..3256fdd1fd 100644
--- a/lightllm/utils/kv_cache_utils.py
+++ b/lightllm/utils/kv_cache_utils.py
@@ -19,13 +19,11 @@
 from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class
 from lightllm.common.kv_cache_mem_manager import (
     MemoryManager,
-    INT8KVMemoryManager,
     CalibrationFP8KVMemoryManager,
     ExportCalibrationMemoryManager,
     PPLINT8KVMemoryManager,
     PPLINT4KVMemoryManager,
     Deepseek2MemoryManager,
-    Deepseek2FP8KVMemoryManager,
 )
 
 from typing import List, Tuple, Optional
@@ -77,17 +75,6 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta":
             scale_head_dim=0,
             scale_data_type=get_llm_data_type(),
         )
-    elif mem_manager_class is Deepseek2FP8KVMemoryManager:
-        cpu_cache_meta = CpuKVCacheMeta(
-            page_num=0,
-            token_page_size=args.cpu_cache_token_page_size,
-            layer_num=get_layer_num(args.model_dir),
-            num_heads=1,
-            head_dim=512 + 64 + 2,
-            data_type=torch.uint8,
-            scale_head_dim=0,
-            scale_data_type=get_llm_data_type(),
-        )
     elif mem_manager_class is MemoryManager:
         cpu_cache_meta = CpuKVCacheMeta(
             page_num=0,
diff --git a/test/acc/test_deepseekr1.sh b/test/acc/test_deepseekr1.sh
index e167303a35..180d2d4e20 100644
--- a/test/acc/test_deepseekr1.sh
+++ b/test/acc/test_deepseekr1.sh
@@ -1,4 +1,4 @@
-LOADWORKER=18 python -m lightllm.server.api_server --model_dir /mtc/models/DeepSeek-R1 --tp 8 --port 8089 --enable_fa3
+LOADWORKER=18 python -m lightllm.server.api_server --batch_max_tokens 6000 --model_dir /mtc/models/DeepSeek-R1 --tp 8 --port 8089 --llm_prefill_att_backend fa3 --llm_decode_att_backend fa3
 
 
 
diff --git a/test/acc/test_deepseekr1_mtp.sh b/test/acc/test_deepseekr1_mtp.sh
index 046314a728..7eaffd4993 100644
--- a/test/acc/test_deepseekr1_mtp.sh
+++ b/test/acc/test_deepseekr1_mtp.sh
@@ -1,3 +1,3 @@
-LOADWORKER=18 python -m lightllm.server.api_server --model_dir /mtc/models/DeepSeek-R1 --tp 8 --port 8089 --mem_fraction 0.75 --enable_fa3 --batch_max_tokens 6000 --mtp_mode eagle_with_att --mtp_draft_model_dir /mtc/models/DeepSeek-R1-NextN --mtp_step 2
+LOADWORKER=18 python -m lightllm.server.api_server --model_dir /mtc/models/DeepSeek-R1 --tp 8 --port 8089 --mem_fraction 0.75 --llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 --batch_max_tokens 6000 --mtp_mode eagle_with_att --mtp_draft_model_dir /mtc/models/DeepSeek-R1-NextN --mtp_step 2
 
 HF_ALLOW_CODE_EVAL=1 HF_DATASETS_OFFLINE=0 lm_eval --model local-completions --model_args '{"model":"deepseek-ai/DeepSeek-R1", "base_url":"http://localhost:8089/v1/completions", "max_length": 16384}' --tasks gsm8k --batch_size 500 --confirm_run_unsafe_code
\ No newline at end of file
diff --git a/test/acc/test_deepseekr1_mtp_ep.sh b/test/acc/test_deepseekr1_mtp_ep.sh
index 2ea5f74387..0467f76e6a 100644
--- a/test/acc/test_deepseekr1_mtp_ep.sh
+++ b/test/acc/test_deepseekr1_mtp_ep.sh
@@ -1,3 +1,3 @@
-LOADWORKER=18 MOE_MODE=EP NUM_MAX_DISPATCH_TOKENS_PER_RANK=256  python -m lightllm.server.api_server --model_dir /mtc/models/DeepSeek-R1 --tp 8 --dp 8 --port 8089 --max_total_token_num 60000 --graph_max_batch_size 16 --enable_fa3 --batch_max_tokens 6000 --mtp_mode eagle_with_att --mtp_draft_model_dir /mtc/models/DeepSeek-R1-NextN --mtp_step 2
+LOADWORKER=18 MOE_MODE=EP NUM_MAX_DISPATCH_TOKENS_PER_RANK=256  python -m lightllm.server.api_server --model_dir /mtc/models/DeepSeek-R1 --tp 8 --dp 8 --port 8089 --max_total_token_num 60000 --graph_max_batch_size 16 --llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 --batch_max_tokens 6000 --mtp_mode eagle_with_att --mtp_draft_model_dir /mtc/models/DeepSeek-R1-NextN --mtp_step 2
 
 HF_ALLOW_CODE_EVAL=1 HF_DATASETS_OFFLINE=0 lm_eval --model local-completions --model_args '{"model":"deepseek-ai/DeepSeek-R1", "base_url":"http://localhost:8089/v1/completions", "max_length": 16384}' --tasks gsm8k --batch_size 32 --confirm_run_unsafe_code
\ No newline at end of file
diff --git a/test/acc/test_qwen2.sh b/test/acc/test_qwen2.sh
index 265d679e8a..bb5603b5be 100644
--- a/test/acc/test_qwen2.sh
+++ b/test/acc/test_qwen2.sh
@@ -1,5 +1,5 @@
 # first
-LOADWORKER=18 CUDA_VISIBLE_DEVICES=6,7 python -m lightllm.server.api_server --model_dir /root/.cache/huggingface/hub/models--Qwen--Qwen2.5-Math-7B-Instruct/snapshots/ef9926d75ab1d54532f6a30dd5e760355eb9aa4d --tp 2 --port 8089 --enable_fa3
+LOADWORKER=18 CUDA_VISIBLE_DEVICES=6,7 python -m lightllm.server.api_server --model_dir /root/.cache/huggingface/hub/models--Qwen--Qwen2.5-Math-7B-Instruct/snapshots/ef9926d75ab1d54532f6a30dd5e760355eb9aa4d --tp 2 --port 8089  --llm_prefill_att_backend fa3 --llm_decode_att_backend fa3
 
 # second
 HF_ALLOW_CODE_EVAL=1 HF_DATASETS_OFFLINE=0 lm_eval --model local-completions --model_args '{"model":"Qwen/Qwen2.5-Math-7B-Instruct", "base_url":"http://localhost:8089/v1/completions", "max_length": 16384}' --tasks gsm8k --batch_size 500 --confirm_run_unsafe_code
\ No newline at end of file
diff --git a/test/acc/test_qwen3.sh b/test/acc/test_qwen3.sh
index c0da5ec96e..36a3c96804 100644
--- a/test/acc/test_qwen3.sh
+++ b/test/acc/test_qwen3.sh
@@ -1,5 +1,5 @@
 # first
-LOADWORKER=18 CUDA_VISIBLE_DEVICES=6,7 python -m lightllm.server.api_server --model_dir /mtc/models/qwen3-8b --tp 2 --port 8089 --enable_fa3
+LOADWORKER=18 CUDA_VISIBLE_DEVICES=6,7 python -m lightllm.server.api_server --model_dir /mtc/models/qwen3-8b --tp 2 --port 8089 --llm_prefill_att_backend fa3 --llm_decode_att_backend fa3
 
 # second
 HF_ALLOW_CODE_EVAL=1 HF_DATASETS_OFFLINE=0 lm_eval --model local-completions --model_args '{"model":"qwen/qwen3-8b", "base_url":"http://localhost:8089/v1/completions", "max_length": 16384}' --tasks gsm8k --batch_size 500 --confirm_run_unsafe_code
\ No newline at end of file
diff --git a/test/benchmark/static_inference/model_infer.py b/test/benchmark/static_inference/model_infer.py
index 3fc7ee4b45..7f1c2b493f 100644
--- a/test/benchmark/static_inference/model_infer.py
+++ b/test/benchmark/static_inference/model_infer.py
@@ -41,7 +41,10 @@ def test_model_inference(args):
             "run_mode": "normal",
             "max_seq_length": args.max_req_total_len,
             "disable_cudagraph": args.disable_cudagraph,
-            "mode": args.mode,
+            "llm_prefill_att_backend": args.llm_prefill_att_backend,
+            "llm_decode_att_backend": args.llm_decode_att_backend,
+            "llm_kv_type": args.llm_kv_type,
+            "llm_kv_quant_group_size": args.llm_kv_quant_group_size,
         }
         proc = multiprocessing.Process(
             target=tppart_model_infer,
@@ -73,7 +76,6 @@ def overlap_prefill(
 ):
     _0_batch_size = batch_size // 2
     _0_total_token_num = total_token_num // 2
-    _0_max_len_in_batch = max_len_in_batch
     _0_input_ids = input_ids[: total_token_num // 2]
     _0_mem_indexes = mem_indexes[: total_token_num // 2]
     _0_b_req_idx = b_req_idx[: batch_size // 2]
@@ -83,7 +85,6 @@ def overlap_prefill(
     micro_batch1 = ModelInput(
         batch_size=_0_batch_size,
         total_token_num=_0_total_token_num,
-        max_len_in_batch=_0_max_len_in_batch,
         input_ids=_0_input_ids,
         b_req_idx=_0_b_req_idx,
         b_mtp_index=_0_b_mtp_index,
@@ -96,7 +97,6 @@ def overlap_prefill(
 
     _1_batch_size = batch_size - batch_size // 2
     _1_total_token_num = total_token_num - total_token_num // 2
-    _1_max_len_in_batch = max_len_in_batch
     _1_input_ids = input_ids[total_token_num // 2 :]
     _1_mem_indexes = mem_indexes[total_token_num // 2 :]
     _1_b_req_idx = b_req_idx[batch_size // 2 :]
@@ -107,7 +107,6 @@ def overlap_prefill(
     micro_batch2 = ModelInput(
         batch_size=_1_batch_size,
         total_token_num=_1_total_token_num,
-        max_len_in_batch=_1_max_len_in_batch,
         input_ids=_1_input_ids,
         b_req_idx=_1_b_req_idx,
         b_mtp_index=_1_b_mtp_index,
@@ -129,7 +128,6 @@ def overlap_decode(
 ):
     _0_batch_size = batch_size // 2
     _0_total_token_num = total_token_num // 2
-    _0_max_len_in_batch = max_len_in_batch
     _0_input_ids = input_ids[: batch_size // 2]
     _0_mem_indexes = mem_indexes[: batch_size // 2]
     _0_b_req_idx = b_req_idx[: batch_size // 2]
@@ -138,7 +136,6 @@ def overlap_decode(
     micro_batch1 = ModelInput(
         batch_size=_0_batch_size,
         total_token_num=_0_total_token_num,
-        max_len_in_batch=_0_max_len_in_batch,
         input_ids=_0_input_ids,
         b_req_idx=_0_b_req_idx,
         b_mtp_index=_0_b_mtp_index,
@@ -149,7 +146,6 @@ def overlap_decode(
 
     _1_batch_size = batch_size - batch_size // 2
     _1_total_token_num = total_token_num - total_token_num // 2
-    _1_max_len_in_batch = max_len_in_batch
     _1_input_ids = input_ids[batch_size // 2 :]
     _1_mem_indexes = mem_indexes[batch_size // 2 :]
     _1_b_req_idx = b_req_idx[batch_size // 2 :]
@@ -159,7 +155,6 @@ def overlap_decode(
     micro_batch2 = ModelInput(
         batch_size=_1_batch_size,
         total_token_num=_1_total_token_num,
-        max_len_in_batch=_1_max_len_in_batch,
         input_ids=_1_input_ids,
         b_req_idx=_1_b_req_idx,
         b_mtp_index=_1_b_mtp_index,
@@ -191,7 +186,6 @@ def prefill(
     model_input = ModelInput(
         batch_size=batch_size,
         total_token_num=total_token_num,
-        max_len_in_batch=max_len_in_batch,
         max_q_seq_len=max_len_in_batch,
         max_kv_seq_len=max_len_in_batch,
         max_cache_len=0,
@@ -217,7 +211,6 @@ def decode(
     model_input = ModelInput(
         batch_size=batch_size,
         total_token_num=total_token_num,
-        max_len_in_batch=max_len_in_batch,
         max_q_seq_len=1,
         max_kv_seq_len=max_len_in_batch,
         input_ids=input_ids,
diff --git a/test/benchmark/static_inference/model_infer_mtp.py b/test/benchmark/static_inference/model_infer_mtp.py
index 942af0f883..07ad52a132 100644
--- a/test/benchmark/static_inference/model_infer_mtp.py
+++ b/test/benchmark/static_inference/model_infer_mtp.py
@@ -129,7 +129,6 @@ def run_forward_once(args, input_len, output_len, batch_size, main_model, draft_
     model_input = ModelInput(
         batch_size=batch_size,
         total_token_num=total_token_num,
-        max_len_in_batch=input_len,
         input_ids=test_data,
         mem_indexes=mem_indexes,
         b_req_idx=b_req_idx,
@@ -197,7 +196,6 @@ def run_forward_once(args, input_len, output_len, batch_size, main_model, draft_
     model_input = ModelInput(
         batch_size=batch_size * (len(draft_models) + 1),
         total_token_num=nopad_total_token_num,
-        max_len_in_batch=nopad_max_len_in_batch,
         input_ids=decode_input_ids,
         mem_indexes=mem_indexes,
         b_req_idx=nopad_b_seq_idx,
diff --git a/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py b/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py
index d391d30650..f32b093448 100644
--- a/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py
+++ b/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py
@@ -4,7 +4,7 @@
 import torch.multiprocessing as mp
 from typing import List
 from lightllm.utils.log_utils import init_logger
-from lightllm.models.llama.triton_kernel.ppl_int8kv_flash_decoding_diverse_stage1 import (
+from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import (
     flash_decode_stage1,
     GQADiverseDecodeStage1KernelConfig,
 )
diff --git a/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py b/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py
new file mode 100644
index 0000000000..13c8945e59
--- /dev/null
+++ b/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py
@@ -0,0 +1,296 @@
+import torch
+import os
+import torch.multiprocessing as mp
+from typing import List
+from lightllm.utils.log_utils import init_logger
+from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import (
+    flash_decode_stage2,
+    GQADiverseDecodeStage2KernelConfig,
+)
+from lightllm.utils.watchdog_utils import Watchdog
+
+logger = init_logger(__name__)
+
+
+def set_seed():
+    import torch
+    import random
+    import numpy as np
+
+    seed = 42
+    torch.manual_seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+    return
+
+
+@torch.no_grad()
+def test_decode_attentions(
+    block_seq: int,
+    batch_size: int,
+    seq_len: int,
+    dtype: torch.dtype,
+    test_count: int = 20,
+    **run_config,
+):
+    set_seed()
+    shared_seq_len = 0
+    num_heads = 32
+    kv_head_num = 8
+    head_dim = 128
+    max_len_in_batch = 8192
+    quant_group_size = 8
+
+    args = []
+    for _ in range(test_count):
+        q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=dtype, device="cuda") / 10
+        kv_shape = (batch_size * seq_len, kv_head_num, head_dim)
+        kv_scale_shape = (batch_size * seq_len, kv_head_num, head_dim // quant_group_size)
+        k = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda")
+        k_scale = torch.ones(size=kv_scale_shape, dtype=dtype, device="cuda")
+        v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda")
+        v_scale = torch.ones(size=kv_scale_shape, dtype=dtype, device="cuda")
+        Req_to_tokens = torch.arange(0, seq_len * batch_size, dtype=torch.int32, device="cuda").view(
+            batch_size, seq_len
+        )
+        B_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda")
+        b_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda")
+        b_shared_seq_len = torch.full((batch_size,), shared_seq_len, dtype=torch.int32, device="cuda")
+        mid_out = torch.zeros(
+            size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2, head_dim),
+            dtype=q.dtype,
+            device="cuda",
+        )
+        mid_out_logsumexp = torch.zeros(
+            size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2),
+            dtype=q.dtype,
+            device="cuda",
+        )
+        arg_list, kwargs = (
+            q,
+            k,
+            k_scale,
+            v,
+            v_scale,
+            Req_to_tokens,
+            B_req_idx,
+            b_seq_len,
+            b_shared_seq_len,
+            max_len_in_batch,
+            mid_out,
+            mid_out_logsumexp,
+            block_seq,
+        ), dict(run_config=run_config)
+        args.append((arg_list, kwargs))
+
+    graph = torch.cuda.CUDAGraph()
+    arg_list, kwargs = args[0]
+    flash_decode_stage2(*arg_list, **kwargs)
+    with torch.cuda.graph(graph):
+        for index in range(test_count):
+            arg_list, kwargs = args[index]
+            flash_decode_stage2(*arg_list, **kwargs)
+
+    graph.replay()
+
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+    start_event.record()
+    graph.replay()
+    end_event.record()
+    end_event.synchronize()
+
+    cost_time = start_event.elapsed_time(end_event=end_event)
+
+    logger.info(f"bf16 {seq_len} cost time: {cost_time} ms")
+    return cost_time
+
+
+def worker(
+    block_seq: int,
+    batch_size: int,
+    seq_len: int,
+    dtype: torch.dtype,
+    test_count: int,
+    test_configs,
+    queue,
+):
+    dog = Watchdog(timeout=10)
+    dog.start()
+
+    try:
+        for index in range(len(test_configs)):
+            tuning_config = test_configs[index]
+            cost_time = test_decode_attentions(
+                block_seq=block_seq,
+                batch_size=batch_size,
+                seq_len=seq_len,
+                dtype=dtype,
+                test_count=test_count,
+                **tuning_config,
+            )
+            dog.heartbeat()
+            queue.put(cost_time)
+    except Exception as ex:
+        logger.error(str(ex) + f" config {tuning_config} batch_size {batch_size} seq_len {seq_len} dtype {dtype}")
+        import sys
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(-1)
+        pass
+
+
+def get_test_configs(split_id, split_count):
+    index = 0
+    for block_n in [16, 32, 64]:
+        for num_warps in [
+            2,
+            4,
+            8,
+            16,
+        ]:
+            for num_stages in [
+                1,
+                2,
+                3,
+                4,
+                5,
+                7,
+                9,
+                10,
+                11,
+            ]:
+                t_config = {
+                    "BLOCK_N": block_n,
+                    "num_warps": num_warps,
+                    "num_stages": num_stages,
+                }
+                if index % split_count == split_id:
+                    yield t_config
+                    index += 1
+                else:
+                    index += 1
+
+
+def tuning_configs(
+    device_id: int,
+    device_count: int,
+    block_seq: int,
+    batch_size: int,
+    seq_len: int,
+    dtype: torch.dtype,
+    test_count: int,
+):
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
+    best_config, best_cost_time = None, 10000000
+    queue = mp.Queue()
+    test_configs = []
+    for t_config in get_test_configs(device_id, device_count):
+        test_configs.append(t_config)
+        if len(test_configs) < 64:
+            continue
+
+        p = mp.Process(
+            target=worker,
+            args=(
+                block_seq,
+                batch_size,
+                seq_len,
+                dtype,
+                test_count,
+                test_configs,
+                queue,
+            ),
+        )
+        p.start()
+        p.join()
+
+        while len(test_configs) != 0:
+            try:
+                cost_time = queue.get_nowait()
+                logger.info(f"get {test_configs[0]} cost_time: {cost_time}")
+                if cost_time < best_cost_time:
+                    best_config = test_configs[0]
+                    best_cost_time = cost_time
+                    logger.info(f"cur best {best_config}, {best_cost_time}")
+                del test_configs[0:1]
+            except:
+                logger.info(f"cur best {best_config}, {best_cost_time}")
+                del test_configs[0:1]
+                break
+
+    while len(test_configs) != 0:
+        p = mp.Process(
+            target=worker,
+            args=(
+                block_seq,
+                batch_size,
+                seq_len,
+                dtype,
+                test_count,
+                test_configs,
+                queue,
+            ),
+        )
+        p.start()
+        p.join()
+
+        while len(test_configs) != 0:
+            try:
+                cost_time = queue.get_nowait()
+                logger.info(f"get {test_configs[0]} cost_time: {cost_time}")
+                if cost_time < best_cost_time:
+                    best_config = test_configs[0]
+                    best_cost_time = cost_time
+                    logger.info(f"cur best {best_config}, {best_cost_time}")
+                del test_configs[0:1]
+            except:
+                logger.info(f"cur best {best_config}, {best_cost_time}")
+                del test_configs[0:1]
+                break
+
+    logger.info(f"{best_config} best cost: {best_cost_time}")
+    return best_config, best_cost_time
+
+
+if __name__ == "__main__":
+    torch.multiprocessing.set_start_method("spawn")
+
+    from lightllm.utils.tuning_utils import mp_tuning
+    import collections
+
+    block_seq = 256
+    batch_sizes = [8, 16, 32, 64]
+    seq_lens = [32, 64, 128, 256]
+    num_heads = 32
+    kv_head_num = 8
+    q_head_dim = 128
+    gqa_group_size = num_heads // kv_head_num
+
+    store_json_ans = collections.defaultdict(dict)
+
+    for seq_len in seq_lens:
+        for batch_size in batch_sizes:
+            ans = mp_tuning(
+                tuning_configs,
+                {
+                    "block_seq": block_seq,
+                    "batch_size": batch_size,
+                    "seq_len": seq_len,
+                    "dtype": torch.bfloat16,
+                    "test_count": 1,
+                },
+            )
+            store_json_ans[seq_len][batch_size] = ans
+
+            GQADiverseDecodeStage2KernelConfig.save_config(
+                gqa_group_size=gqa_group_size,
+                q_head_dim=q_head_dim,
+                block_seq=block_seq,
+                out_dtype=str(torch.bfloat16),
+                config_json=store_json_ans,
+            )
diff --git a/test/start_scripts/README.md b/test/start_scripts/README.md
index f5dae19b92..8ed44a2753 100644
--- a/test/start_scripts/README.md
+++ b/test/start_scripts/README.md
@@ -99,7 +99,6 @@ sh multi_pd_master/pd_decode.sh <host> <config_server_host>
 ### Environment Variables
 
 - `LOADWORKER`: Model loading thread count, recommended 8-18
-- `MOE_MODE`: Expert parallelism mode, set to EP to enable expert parallelism
 - `DISABLE_KV_TRANS_USE_P2P`: Disable P2P communication optimization to transfer kv data
 - `CUDA_VISIBLE_DEVICES`: Specify GPU devices to use
 
@@ -108,7 +107,7 @@ sh multi_pd_master/pd_decode.sh <host> <config_server_host>
 - `--model_dir`: Model file path
 - `--tp`: Tensor parallelism degree
 - `--dp`: Data parallelism degree
-- `--enable_fa3`: Enable Flash Attention 3.0
+- `--enable_ep_mode`: enable expert parallel 
 - `--nnodes`: Total number of nodes
 - `--node_rank`: Current node rank
 - `--nccl_host`: NCCL communication host address
diff --git a/test/start_scripts/draft.sh b/test/start_scripts/draft.sh
index 866f5f2fa5..04f573cd6d 100644
--- a/test/start_scripts/draft.sh
+++ b/test/start_scripts/draft.sh
@@ -3,7 +3,7 @@ LOADWORKER=18 python -m lightllm.server.api_server \
 --model_dir /mtc/models/qwen3-8b --tp 2 --dp 1 --enable_cpu_cache  --cpu_cache_storage_size 66 --cpu_cache_token_page_size 128 \
 --batch_max_tokens 4096 --chunked_prefill_size 2048 \
 --max_total_token_num 20000 \
---mode "ppl_int8kv_flashdecoding" | tee log.txt
+--llm_kv_type int8kv | tee log.txt
 
 
 # 精度评测命令
@@ -16,7 +16,7 @@ HF_ALLOW_CODE_EVAL=1 HF_DATASETS_OFFLINE=0 lm_eval --model local-completions \
 LOADWORKER=18 python -m lightllm.server.api_server \
 --model_dir /mtc/DeepSeek-R1 \
 --tp 8 \
---enable_fa3 \
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 \
 --batch_max_tokens 4096 --chunked_prefill_size 2048 \
 --max_total_token_num 20000 \
 --enable_cpu_cache  --cpu_cache_storage_size 66 --cpu_cache_token_page_size 128
diff --git a/test/start_scripts/multi_node_ep_node0.sh b/test/start_scripts/multi_node_ep_node0.sh
index 3a139968a6..2e46f70ac2 100644
--- a/test/start_scripts/multi_node_ep_node0.sh
+++ b/test/start_scripts/multi_node_ep_node0.sh
@@ -2,15 +2,15 @@
 # nccl_host: the ip of the nccl host
 # sh multi_node_ep_node0.sh <nccl_host>
 export nccl_host=$1
-MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
+LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
 --model_dir /path/DeepSeek-R1 \
 --tp 16 \
 --dp 16 \
---enable_fa3 \
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 \
 --nnodes 2 \
 --node_rank 0 \
 --nccl_host $nccl_host \
---nccl_port 2732 
+--nccl_port 2732 --enable_ep_moe
 # if you want to enable microbatch overlap, you can uncomment the following lines
 #--enable_prefill_microbatch_overlap
 #--enable_decode_microbatch_overlap
\ No newline at end of file
diff --git a/test/start_scripts/multi_node_ep_node1.sh b/test/start_scripts/multi_node_ep_node1.sh
index b24a598688..a8ce02a8b6 100644
--- a/test/start_scripts/multi_node_ep_node1.sh
+++ b/test/start_scripts/multi_node_ep_node1.sh
@@ -2,15 +2,15 @@
 # nccl_host: the ip of the nccl host
 # sh multi_node_ep_node1.sh <nccl_host>
 export nccl_host=$1
-MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
+LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
 --model_dir /path/DeepSeek-R1 \
 --tp 16 \
 --dp 16 \
---enable_fa3 \
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 \
 --nnodes 2 \
 --node_rank 1 \
 --nccl_host $nccl_host \
---nccl_port 2732 
+--nccl_port 2732 --enable_ep_moe
 # if you want to enable microbatch overlap, you can uncomment the following lines
 #--enable_prefill_microbatch_overlap
 #--enable_decode_microbatch_overlap
\ No newline at end of file
diff --git a/test/start_scripts/multi_node_tp_node0.sh b/test/start_scripts/multi_node_tp_node0.sh
index b86bdeb358..d750da93ca 100644
--- a/test/start_scripts/multi_node_tp_node0.sh
+++ b/test/start_scripts/multi_node_tp_node0.sh
@@ -5,7 +5,7 @@ export nccl_host=$1
 LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
 --model_dir /path/DeepSeek-R1 \
 --tp 16 \
---enable_fa3 \
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 \
 --nnodes 2 \
 --node_rank 0 \
 --nccl_host $nccl_host \
diff --git a/test/start_scripts/multi_node_tp_node1.sh b/test/start_scripts/multi_node_tp_node1.sh
index 378977ab2e..cb495496e8 100644
--- a/test/start_scripts/multi_node_tp_node1.sh
+++ b/test/start_scripts/multi_node_tp_node1.sh
@@ -5,7 +5,7 @@ export nccl_host=$1
 LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
 --model_dir /path/DeepSeek-R1 \
 --tp 16 \
---enable_fa3 \
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 \
 --nnodes 2 \
 --node_rank 1 \
 --nccl_host $nccl_host \
diff --git a/test/start_scripts/multi_pd_master/pd_decode.sh b/test/start_scripts/multi_pd_master/pd_decode.sh
index 4cefef6fb2..2b7bb80d76 100644
--- a/test/start_scripts/multi_pd_master/pd_decode.sh
+++ b/test/start_scripts/multi_pd_master/pd_decode.sh
@@ -13,7 +13,7 @@ MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
 --nccl_port 12322 \
 --tp 8 \
 --dp 8 \
---enable_fa3 \
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3
 --config_server_host $config_server_host \
 --config_server_port 60088
 # if you want to enable microbatch overlap, you can uncomment the following lines
diff --git a/test/start_scripts/multi_pd_master/pd_prefill.sh b/test/start_scripts/multi_pd_master/pd_prefill.sh
index b845da435d..4087c16557 100644
--- a/test/start_scripts/multi_pd_master/pd_prefill.sh
+++ b/test/start_scripts/multi_pd_master/pd_prefill.sh
@@ -5,7 +5,7 @@
 export host=$1
 export config_server_host=$2
 nvidia-cuda-mps-control -d
-MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+LOADWORKER=18 python -m lightllm.server.api_server \
 --model_dir /path/DeepSeek-R1 \
 --run_mode "prefill" \
 --host $host \
@@ -13,9 +13,10 @@ MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
 --tp 8 \
 --dp 8 \
 --nccl_port 2732 \
---enable_fa3 \
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 \
 --disable_cudagraph \
 --config_server_host $config_server_host \
---config_server_port 60088
+--config_server_port 60088 \
+--enable_ep_moe
 # if you want to enable microbatch overlap, you can uncomment the following lines
 #--enable_prefill_microbatch_overlap
\ No newline at end of file
diff --git a/test/start_scripts/single_node_ep.sh b/test/start_scripts/single_node_ep.sh
index cad172d515..05798a2049 100644
--- a/test/start_scripts/single_node_ep.sh
+++ b/test/start_scripts/single_node_ep.sh
@@ -1,9 +1,10 @@
 # H200 single node deepseek R1 dpep mode
-MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
+LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
 --model_dir /path/DeepSeek-R1 \
 --tp 8 \
 --dp 8 \
---enable_fa3
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 \
+--enable_ep_moe
 # if you want to enable microbatch overlap, you can uncomment the following lines
 #--enable_prefill_microbatch_overlap \
 #--enable_decode_microbatch_overlap \
diff --git a/test/start_scripts/single_node_tp.sh b/test/start_scripts/single_node_tp.sh
index 1fb461bb11..ee10b6c101 100644
--- a/test/start_scripts/single_node_tp.sh
+++ b/test/start_scripts/single_node_tp.sh
@@ -2,7 +2,7 @@
 LOADWORKER=18 python -m lightllm.server.api_server --port 8088 \
 --model_dir /path/DeepSeek-R1 \
 --tp 8 \
---enable_fa3
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3
 # if you want to enable microbatch overlap, you can uncomment the following lines
 #--enable_prefill_microbatch_overlap \
 #--enable_decode_microbatch_overlap \
diff --git a/test/start_scripts/single_node_tp_cpu_cache_enable.sh b/test/start_scripts/single_node_tp_cpu_cache_enable.sh
index 3caabb59bd..47da83dbe9 100644
--- a/test/start_scripts/single_node_tp_cpu_cache_enable.sh
+++ b/test/start_scripts/single_node_tp_cpu_cache_enable.sh
@@ -3,7 +3,7 @@ LOADWORKER=18 python -m lightllm.server.api_server \
 --model_dir /mtc/models/qwen3-8b --tp 2 --dp 1 --enable_cpu_cache  --cpu_cache_storage_size 66 --cpu_cache_token_page_size 128 \
 --batch_max_tokens 4096 --chunked_prefill_size 2048 \
 --max_total_token_num 20000 \
---mode "ppl_int8kv_flashdecoding" | tee log.txt
+--llm_kv_type int8kv | tee log.txt
 
 
 # 精度评测命令
diff --git a/test/start_scripts/single_pd_master/pd_decode.sh b/test/start_scripts/single_pd_master/pd_decode.sh
index ae16b96ad4..5471f1d84d 100644
--- a/test/start_scripts/single_pd_master/pd_decode.sh
+++ b/test/start_scripts/single_pd_master/pd_decode.sh
@@ -5,7 +5,7 @@
 export host=$1
 export pd_master_ip=$2
 nvidia-cuda-mps-control -d
-MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+LOADWORKER=18 python -m lightllm.server.api_server \
 --model_dir /path/DeepSeek-R1 \
 --run_mode "decode" \
 --tp 8 \
@@ -13,7 +13,8 @@ MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
 --host $host \
 --port 8121 \
 --nccl_port 12322 \
---enable_fa3 \
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 \
+--enable_ep_moe \
 --pd_master_ip $pd_master_ip \
 --pd_master_port 60011 
 # if you want to enable microbatch overlap, you can uncomment the following lines
diff --git a/test/start_scripts/single_pd_master/pd_nixl_decode.sh b/test/start_scripts/single_pd_master/pd_nixl_decode.sh
index 1b43c11cc4..9354661d3e 100644
--- a/test/start_scripts/single_pd_master/pd_nixl_decode.sh
+++ b/test/start_scripts/single_pd_master/pd_nixl_decode.sh
@@ -10,7 +10,7 @@ export UCX_LOG_LEVEL=info
 export UCX_TLS=rc,cuda,gdr_copy
 
 nvidia-cuda-mps-control -d
-MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+LOADWORKER=18 python -m lightllm.server.api_server \
 --model_dir /path/DeepSeek-R1 \
 --run_mode "nixl_decode" \
 --tp 8 \
@@ -18,7 +18,8 @@ MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
 --host $host \
 --port 8121 \
 --nccl_port 12322 \
---enable_fa3 \
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 \
+--enable_ep_moe \
 --pd_master_ip $pd_master_ip \
 --pd_master_port 60011 
 # if you want to enable microbatch overlap, you can uncomment the following lines
diff --git a/test/start_scripts/single_pd_master/pd_nixl_prefill.sh b/test/start_scripts/single_pd_master/pd_nixl_prefill.sh
index 303de29758..ea3188be39 100644
--- a/test/start_scripts/single_pd_master/pd_nixl_prefill.sh
+++ b/test/start_scripts/single_pd_master/pd_nixl_prefill.sh
@@ -11,7 +11,7 @@ export UCX_TLS=rc,cuda,gdr_copy
 export host=$1
 export pd_master_ip=$2
 nvidia-cuda-mps-control -d 
-MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+LOADWORKER=18 python -m lightllm.server.api_server \
 --model_dir /path/DeepSeek-R1 \
 --run_mode "nixl_prefill" \
 --tp 8 \
@@ -19,7 +19,8 @@ MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
 --host $host \
 --port 8019 \
 --nccl_port 2732 \
---enable_fa3 \
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 \
+--enable_ep_moe \
 --disable_cudagraph \
 --pd_master_ip $pd_master_ip \
 --pd_master_port 60011 
diff --git a/test/start_scripts/single_pd_master/pd_prefill.sh b/test/start_scripts/single_pd_master/pd_prefill.sh
index f6e2e4b685..131b3754e6 100644
--- a/test/start_scripts/single_pd_master/pd_prefill.sh
+++ b/test/start_scripts/single_pd_master/pd_prefill.sh
@@ -5,7 +5,7 @@
 export host=$1
 export pd_master_ip=$2
 nvidia-cuda-mps-control -d 
-MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
+LOADWORKER=18 python -m lightllm.server.api_server \
 --model_dir /path/DeepSeek-R1 \
 --run_mode "prefill" \
 --tp 8 \
@@ -13,9 +13,10 @@ MOE_MODE=EP LOADWORKER=18 python -m lightllm.server.api_server \
 --host $host \
 --port 8019 \
 --nccl_port 2732 \
---enable_fa3 \
+--llm_prefill_att_backend fa3 --llm_decode_att_backend fa3 \
 --disable_cudagraph \
 --pd_master_ip $pd_master_ip \
---pd_master_port 60011 
+--pd_master_port 60011 \
+--enable_ep_moe
 # if you want to enable microbatch overlap, you can uncomment the following lines
 #--enable_prefill_microbatch_overlap
\ No newline at end of file
diff --git a/test/test_api/test_generate_api.py b/test/test_api/test_generate_api.py
index 05fbda44ea..4ea74b7f6d 100644
--- a/test/test_api/test_generate_api.py
+++ b/test/test_api/test_generate_api.py
@@ -19,7 +19,7 @@ def run(self):
             print("Error:", response.status_code, response.text)
 
 
-url = "http://localhost:8000/generate"
+url = "http://localhost:8089/generate"
 headers = {"Content-Type": "application/json"}
 
 for i in range(1):
diff --git a/unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py
similarity index 81%
rename from unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse.py
rename to unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py
index 67463fa7ad..a01bbf32d8 100644
--- a/unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse.py
+++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py
@@ -1,4 +1,5 @@
 import pytest
+
 import torch
 from lightllm.utils.light_utils import light_ops
 
@@ -21,7 +22,7 @@ class MockInferState:
     def __init__(
         self,
         batch_size,
-        max_len_in_batch,
+        max_kv_seq_len,
         req_to_tokens,
         b_req_idx,
         b_seq_len,
@@ -29,7 +30,7 @@ def __init__(
         b_mark_shared_group=None,
     ):
         self.batch_size = batch_size
-        self.max_len_in_batch = max_len_in_batch
+        self.max_kv_seq_len = max_kv_seq_len
         self.req_manager = MockReqManager(req_to_tokens)
         self.b_req_idx = b_req_idx
         self.b_seq_len = b_seq_len
@@ -39,30 +40,32 @@ def __init__(
 
 # @pytest.mark.parametrize("shared_seq_len", [512])
 @pytest.mark.parametrize("shared_seq_len", [0, 77, 256, 311, 512, 550])
-def test_token_decode_attention_flash_decoding_diverse_vs_baseline(shared_seq_len):
+@pytest.mark.parametrize("batch_size", list(range(6, 121, 6)))
+def test_token_decode_attention_flash_decoding_diverse_vs_baseline(shared_seq_len, batch_size):
     """
-    测试 ppl_int8kv_flash_decoding_diverse 的 token_decode_attention_flash_decoding
+    测试 int8kv_flash_decoding_diverse 的 token_decode_attention_flash_decoding
     与 ppl_int8kv_flash_decoding (baseline) 的对比。
     """
-    from lightllm.models.llama.triton_kernel.ppl_int8kv_flash_decoding_diverse import (
+
+    from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse import (
         token_decode_attention_flash_decoding as diverse_attention,
     )
-    from lightllm.models.llama.triton_kernel.ppl_int8kv_flash_decoding import (
+    from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.ppl_int8kv_flash_decoding import (
         token_decode_attention_flash_decoding as baseline_attention,
     )
 
-    batch_size = 6
     num_heads = 32
     kv_head_num = 8
     mark_shared_group_size = 3
-    seq_len = 1024
+    seq_len = 3547
     head_dim = 128
     quant_group_size = 8
+    max_len_in_batch = 8192
     test_dtype = torch.bfloat16
 
     # 创建测试数据
-    kv_shape = (batch_size * seq_len, kv_head_num, head_dim)
-    kv_scale_shape = (batch_size * seq_len, kv_head_num, head_dim // quant_group_size)
+    kv_shape = (batch_size * max_len_in_batch, kv_head_num, head_dim)
+    kv_scale_shape = (batch_size * max_len_in_batch, kv_head_num, head_dim // quant_group_size)
 
     q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device="cuda")
 
@@ -73,7 +76,9 @@ def test_token_decode_attention_flash_decoding_diverse_vs_baseline(shared_seq_le
     cache_v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda")
     cache_v_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") / 100.0
 
-    req_to_tokens = torch.arange(0, seq_len * batch_size, dtype=torch.int32, device="cuda").view(batch_size, seq_len)
+    req_to_tokens = torch.arange(0, max_len_in_batch * batch_size, dtype=torch.int32, device="cuda").view(
+        batch_size, max_len_in_batch
+    )
     for i in range(batch_size):
         if i % mark_shared_group_size != 0:
             req_to_tokens[i, :shared_seq_len] = req_to_tokens[i - 1, :shared_seq_len]
@@ -87,7 +92,7 @@ def test_token_decode_attention_flash_decoding_diverse_vs_baseline(shared_seq_le
     # 创建 baseline 的 infer_state (不需要 b_shared_seq_len)
     baseline_infer_state = MockInferState(
         batch_size=batch_size,
-        max_len_in_batch=seq_len,
+        max_kv_seq_len=max_len_in_batch,
         req_to_tokens=req_to_tokens,
         b_req_idx=b_req_idx,
         b_seq_len=b_seq_len,
@@ -96,7 +101,7 @@ def test_token_decode_attention_flash_decoding_diverse_vs_baseline(shared_seq_le
     # 创建 diverse 的 infer_state
     diverse_infer_state = MockInferState(
         batch_size=batch_size,
-        max_len_in_batch=seq_len,
+        max_kv_seq_len=max_len_in_batch,
         req_to_tokens=req_to_tokens,
         b_req_idx=b_req_idx,
         b_seq_len=b_seq_len,
@@ -108,8 +113,6 @@ def test_token_decode_attention_flash_decoding_diverse_vs_baseline(shared_seq_le
     baseline_out = baseline_attention(
         q=q.clone(),
         infer_state=baseline_infer_state,
-        q_head_num=num_heads,
-        head_dim=head_dim,
         cache_k=cache_k,
         cache_k_scale=cache_k_scale,
         cache_v=cache_v,
@@ -120,8 +123,6 @@ def test_token_decode_attention_flash_decoding_diverse_vs_baseline(shared_seq_le
     diverse_out = diverse_attention(
         q=q.clone(),
         infer_state=diverse_infer_state,
-        q_head_num=num_heads,
-        head_dim=head_dim,
         cache_k=cache_k,
         cache_k_scale=cache_k_scale,
         cache_v=cache_v,
@@ -129,7 +130,7 @@ def test_token_decode_attention_flash_decoding_diverse_vs_baseline(shared_seq_le
         alloc_tensor_func=alloc_tensor_func,
     )
 
-    print(f"\nshared_seq_len={shared_seq_len}")
+    print(f"\nshared_seq_len={shared_seq_len}\nbatch_size={batch_size}")
     print(f"baseline_out: {baseline_out[0, 0, :4]}")
     print(f"diverse_out: {diverse_out[0, 0, :4]}")
     print(f"max diff: {(baseline_out - diverse_out).abs().max()}")
diff --git a/unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse_stage1.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py
similarity index 50%
rename from unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse_stage1.py
rename to unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py
index 30e83b88b6..f3cb8de463 100644
--- a/unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse_stage1.py
+++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py
@@ -1,39 +1,48 @@
 import pytest
 import torch
-from lightllm.models.llama.triton_kernel.ppl_int8kv_flash_decoding_diverse_stage1 import flash_decode_stage1
+from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage1 import (
+    flash_decode_stage1,
+)
 
 
-@pytest.fixture
-def setup_tensors():
-    batch_size = 4
-    num_heads = 4
-    kv_head_num = 1
-    seq_len = 256
+def create_tensors(
+    batch_size=4,
+    num_heads=4,
+    kv_head_num=1,
+    seq_len=256,
+    max_len_in_batch=8192,
+    max_batch_group_size=4,
+    kv_len=None,
+    req_to_tokens_len=None,
+):
     head_dim = 128
-    max_len_in_batch = seq_len
     block_seq = 256
-    max_batch_group_size = 4
     quant_group_size = 8
 
     test_dtype = torch.bfloat16
 
-    kv_shape = (batch_size * seq_len, kv_head_num, head_dim)
-    kv_scale_shape = (batch_size * seq_len, kv_head_num, head_dim // quant_group_size)
+    kv_len = max_len_in_batch if kv_len is None else kv_len
+    req_to_tokens_len = max_len_in_batch if req_to_tokens_len is None else req_to_tokens_len
+
+    kv_shape = (batch_size * kv_len, kv_head_num, head_dim)
+    kv_scale_shape = (batch_size * kv_len, kv_head_num, head_dim // quant_group_size)
 
     q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device="cuda")
     k = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda")
     k_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda")
     v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda")
     v_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda")
-    Req_to_tokens = torch.arange(0, seq_len * batch_size, dtype=torch.int32, device="cuda").view(batch_size, seq_len)
+    Req_to_tokens = torch.arange(0, req_to_tokens_len * batch_size, dtype=torch.int32, device="cuda").view(
+        batch_size, req_to_tokens_len
+    )
     B_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda")
     b_shared_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda")
     b_mark_shared_group = torch.ones(batch_size, dtype=torch.int32, device="cuda")
     mid_out = torch.zeros(
-        size=(batch_size, num_heads, (seq_len // block_seq) + 2, head_dim), dtype=q.dtype, device="cuda"
+        size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2, head_dim), dtype=q.dtype, device="cuda"
     )
     mid_out_logsumexp = torch.zeros(
-        size=(batch_size, num_heads, (seq_len // block_seq) + 2), dtype=q.dtype, device="cuda"
+        size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2), dtype=q.dtype, device="cuda"
     )
 
     return {
@@ -54,6 +63,11 @@ def setup_tensors():
     }
 
 
+@pytest.fixture
+def setup_tensors():
+    return create_tensors()
+
+
 def test_flash_decode_stage1_execution(setup_tensors):
     flash_decode_stage1(
         q=setup_tensors["q"],
@@ -81,7 +95,7 @@ def test_flash_decode_stage1_execution(setup_tensors):
     new_k = k.to(q.dtype)
     new_v = v.to(q.dtype)
 
-    from lightllm.models.llama.triton_kernel.gqa_flash_decoding_stage1 import (
+    from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import (
         flash_decode_stage1 as gqa_flash_decode_stage1,
     )
 
@@ -104,3 +118,71 @@ def test_flash_decode_stage1_execution(setup_tensors):
     assert torch.allclose(
         setup_tensors["mid_out_logsumexp"], true_mid_out_logsumexp, atol=1e-2
     ), "LogSumExp output does not match expected values"
+
+
+def autotune_and_benchmark():
+    import triton
+
+    batch_sizes = [8, 16, 32, 64]
+    seq_lens = [1024, 2048, 4096]
+
+    results = []
+    for batch in batch_sizes:
+        for seq in seq_lens:
+            # Clear GPU cache to reduce CUDA Graph capture failures.
+            torch.cuda.empty_cache()
+
+            setup_tensors = create_tensors(
+                batch_size=batch,
+                num_heads=32,
+                kv_head_num=8,
+                seq_len=seq,
+                max_len_in_batch=8192,
+                max_batch_group_size=8,
+                kv_len=seq,
+                req_to_tokens_len=seq,
+            )
+
+            def fn_triton(st=setup_tensors):
+                return flash_decode_stage1(
+                    q=st["q"],
+                    k=st["k"],
+                    k_scale=st["k_scale"],
+                    v=st["v"],
+                    v_scale=st["v_scale"],
+                    Req_to_tokens=st["Req_to_tokens"],
+                    B_req_idx=st["B_req_idx"],
+                    b_shared_seq_len=st["b_shared_seq_len"],
+                    b_mark_shared_group=st["b_mark_shared_group"],
+                    max_len_in_batch=st["max_len_in_batch"],
+                    mid_out=st["mid_out"],
+                    mid_out_logsumexp=st["mid_out_logsumexp"],
+                    block_seq=st["block_seq"],
+                    max_batch_group_size=st["max_batch_group_size"],
+                )
+
+            ms_triton = triton.testing.do_bench_cudagraph(fn_triton, rep=100)
+
+            results.append(
+                {
+                    "batch_size": batch,
+                    "seq_len": seq,
+                    "triton_ms": ms_triton,
+                }
+            )
+            print(results[-1])
+
+            del setup_tensors
+
+    print(f"\n{'='*80}")
+    print("SUMMARY - Performance Comparison")
+    print(f"{'='*80}")
+    print(f"{'batch_size':<8} {'seq_len':<12} {'triton_ms':<12}")
+    print(f"{'-'*80}")
+    for r in results:
+        print(f"{r['batch_size']:<8} {r['seq_len']:<12} {r['triton_ms']:<12.3f}")
+    print(f"{'='*80}")
+
+
+if __name__ == "__main__":
+    autotune_and_benchmark()
diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py
new file mode 100644
index 0000000000..c7d4442543
--- /dev/null
+++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py
@@ -0,0 +1,293 @@
+import pytest
+import torch
+from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import (
+    flash_decode_stage2,
+)
+
+
+def create_tensors(
+    shared_seq_len,
+    batch_size=4,
+    seq_len=256,
+    max_len_in_batch=8192,
+    max_batch_group_size=4,
+    kv_len=None,
+    req_to_tokens_len=None,
+):
+    num_heads = 32
+    kv_head_num = 8
+    head_dim = 128
+    block_seq = 256
+    quant_group_size = 8
+
+    test_dtype = torch.bfloat16
+
+    kv_len = max_len_in_batch if kv_len is None else kv_len
+    req_to_tokens_len = max_len_in_batch if req_to_tokens_len is None else req_to_tokens_len
+
+    kv_shape = (batch_size * kv_len, kv_head_num, head_dim)
+    kv_scale_shape = (batch_size * kv_len, kv_head_num, head_dim // quant_group_size)
+
+    q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device="cuda")
+    k = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda")
+    k_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda")
+    v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda")
+    v_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda")
+    Req_to_tokens = torch.arange(0, req_to_tokens_len * batch_size, dtype=torch.int32, device="cuda").view(
+        batch_size, req_to_tokens_len
+    )
+    B_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda")
+    b_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda")
+    b_shared_seq_len = torch.full((batch_size,), shared_seq_len, dtype=torch.int32, device="cuda")
+    b_mark_shared_group = torch.ones(batch_size, dtype=torch.int32, device="cuda")
+    mid_out = torch.zeros(
+        size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2, head_dim), dtype=q.dtype, device="cuda"
+    )
+    mid_out_logsumexp = torch.zeros(
+        size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2), dtype=q.dtype, device="cuda"
+    )
+
+    return {
+        "q": q,
+        "k": k,
+        "k_scale": k_scale,
+        "v": v,
+        "v_scale": v_scale,
+        "Req_to_tokens": Req_to_tokens,
+        "B_req_idx": B_req_idx,
+        "b_seq_len": b_seq_len,
+        "b_shared_seq_len": b_shared_seq_len,
+        "b_mark_shared_group": b_mark_shared_group,
+        "max_len_in_batch": max_len_in_batch,
+        "mid_out": mid_out,
+        "mid_out_logsumexp": mid_out_logsumexp,
+        "block_seq": block_seq,
+        "max_batch_group_size": max_batch_group_size,
+        "head_dim": head_dim,
+    }
+
+
+@pytest.mark.parametrize("shared_seq_len", [0, 47, 77, 128, 200, 255])
+def test_flash_decode_stage2_execution(shared_seq_len):
+    setup_tensors = create_tensors(shared_seq_len)
+
+    flash_decode_stage2(
+        q=setup_tensors["q"],
+        k=setup_tensors["k"],
+        k_scale=setup_tensors["k_scale"],
+        v=setup_tensors["v"],
+        v_scale=setup_tensors["v_scale"],
+        Req_to_tokens=setup_tensors["Req_to_tokens"],
+        B_req_idx=setup_tensors["B_req_idx"],
+        B_Seqlen=setup_tensors["b_seq_len"],
+        b_shared_seq_len=setup_tensors["b_shared_seq_len"],
+        max_len_in_batch=setup_tensors["max_len_in_batch"],
+        mid_out=setup_tensors["mid_out"],
+        mid_out_logsumexp=setup_tensors["mid_out_logsumexp"],
+        block_seq=setup_tensors["block_seq"],
+    )
+    seq_block_idx = (setup_tensors["b_shared_seq_len"][0].item() + setup_tensors["block_seq"] - 1) // setup_tensors[
+        "block_seq"
+    ]
+    mid_out = setup_tensors["mid_out"][:, :, seq_block_idx:, :]
+    mid_out_logsumexp = setup_tensors["mid_out_logsumexp"][:, :, seq_block_idx:]
+
+    q = setup_tensors["q"]
+    k = setup_tensors["k"]
+    v = setup_tensors["v"]
+    true_mid_out = torch.zeros_like(mid_out)
+    true_mid_out_logsumexp = torch.zeros_like(mid_out_logsumexp)
+    new_q = q
+    new_k = k.to(q.dtype)
+    new_v = v.to(q.dtype)
+
+    b_seq_len = setup_tensors["b_seq_len"] - setup_tensors["b_shared_seq_len"]
+    req_to_tokens = setup_tensors["Req_to_tokens"][:, setup_tensors["b_shared_seq_len"][0].item() :]
+
+    from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import (
+        flash_decode_stage1 as gqa_flash_decode_stage1,
+    )
+
+    gqa_flash_decode_stage1(
+        q=new_q,
+        k=new_k,
+        v=new_v,
+        Req_to_tokens=req_to_tokens,
+        B_req_idx=setup_tensors["B_req_idx"],
+        B_Seqlen=b_seq_len,
+        max_len_in_batch=setup_tensors["max_len_in_batch"],
+        mid_out=true_mid_out,
+        mid_out_logsumexp=true_mid_out_logsumexp,
+        block_seq=setup_tensors["block_seq"],
+    )
+    print(f"\nshared_seq_len={shared_seq_len}")
+    print(f"mid_out: {mid_out[0:4, 0, 0, 0]}")
+    print(f"true_mid_out: {true_mid_out[0:4, 0, 0, 0]}")
+    abs_diff = (mid_out - true_mid_out).abs()
+    max_diff = abs_diff.max()
+    max_diff_idx = abs_diff.argmax()
+    max_diff_idx_unraveled = torch.unravel_index(max_diff_idx, abs_diff.shape)
+    mid_out_value = mid_out[max_diff_idx_unraveled]
+    true_mid_out_value = true_mid_out[max_diff_idx_unraveled]
+    print(f"max abs diff: {max_diff}, mid_out value: {mid_out_value}, " f"true_mid_out value: {true_mid_out_value}")
+
+    assert torch.allclose(
+        mid_out[0:4, 0, 0, 0], true_mid_out[0:4, 0, 0, 0], atol=1e-2
+    ), f"Mid output does not match expected values for shared_seq_len={shared_seq_len}"
+    assert torch.allclose(
+        mid_out_logsumexp, true_mid_out_logsumexp, atol=1e-2
+    ), f"LogSumExp output does not match expected values for shared_seq_len={shared_seq_len}"
+
+
+if __name__ == "__main__":
+    import importlib
+    import triton
+    from lightllm.utils.light_utils import light_ops
+
+    batch_sizes = [8, 16, 32, 64]
+    seq_lens = [32, 64, 128, 256]
+
+    results = []
+    for batch in batch_sizes:
+        for seq in seq_lens:
+            # Clear GPU cache to reduce CUDA Graph capture failures.
+            torch.cuda.empty_cache()
+
+            setup_tensors = create_tensors(
+                shared_seq_len=0,
+                batch_size=batch,
+                seq_len=seq,
+                max_len_in_batch=8192,
+                kv_len=seq,
+                req_to_tokens_len=seq,
+            )
+
+            # Outputs for CUDA implementation
+            mid_out_cuda = setup_tensors["mid_out"].clone()
+            mid_out_logsumexp_cuda = setup_tensors["mid_out_logsumexp"].clone()
+
+            # Outputs for Triton implementation
+            mid_out_triton = setup_tensors["mid_out"].clone()
+            mid_out_logsumexp_triton = setup_tensors["mid_out_logsumexp"].clone()
+
+            # Run CUDA to get reference
+            light_ops.group8_int8kv_flashdecoding_diverse_stage2(
+                setup_tensors["block_seq"],
+                mid_out_cuda,
+                mid_out_logsumexp_cuda,
+                1.0 / (setup_tensors["head_dim"] ** 0.5),
+                setup_tensors["q"],
+                setup_tensors["k"],
+                setup_tensors["k_scale"],
+                setup_tensors["v"],
+                setup_tensors["v_scale"],
+                setup_tensors["Req_to_tokens"],
+                setup_tensors["B_req_idx"],
+                setup_tensors["b_seq_len"],
+                setup_tensors["b_shared_seq_len"],
+                setup_tensors["max_len_in_batch"],
+            )
+
+            # Run Triton
+            flash_decode_stage2(
+                q=setup_tensors["q"],
+                k=setup_tensors["k"],
+                k_scale=setup_tensors["k_scale"],
+                v=setup_tensors["v"],
+                v_scale=setup_tensors["v_scale"],
+                Req_to_tokens=setup_tensors["Req_to_tokens"],
+                B_req_idx=setup_tensors["B_req_idx"],
+                B_Seqlen=setup_tensors["b_seq_len"],
+                b_shared_seq_len=setup_tensors["b_shared_seq_len"],
+                max_len_in_batch=setup_tensors["max_len_in_batch"],
+                mid_out=mid_out_triton,
+                mid_out_logsumexp=mid_out_logsumexp_triton,
+                block_seq=setup_tensors["block_seq"],
+            )
+
+            # Compare results
+            diff_mid_out = torch.abs(mid_out_cuda - mid_out_triton)
+            diff_logsumexp = torch.abs(mid_out_logsumexp_cuda - mid_out_logsumexp_triton)
+            max_diff_out = diff_mid_out.max().item()
+            max_diff_logsumexp = diff_logsumexp.max().item()
+            mean_diff_out = diff_mid_out.mean().item()
+            mean_diff_logsumexp = diff_logsumexp.mean().item()
+
+            cos_sim_out = torch.nn.functional.cosine_similarity(
+                mid_out_cuda.flatten(), mid_out_triton.flatten(), dim=0
+            ).item()
+            cos_sim_logsumexp = torch.nn.functional.cosine_similarity(
+                mid_out_logsumexp_cuda.flatten(), mid_out_logsumexp_triton.flatten(), dim=0
+            ).item()
+
+            print(f"\n[batch={batch}, seq={seq}] Consistency check:")
+            print("  mid_out:")
+            print(f"    max_diff: {max_diff_out:.6f}, mean_diff: {mean_diff_out:.6f}, cosine_sim: {cos_sim_out:.8f}")
+            print("  logsumexp:")
+            print(
+                f"    max_diff: {max_diff_logsumexp:.6f}, "
+                f"mean_diff: {mean_diff_logsumexp:.6f}, "
+                f"cosine_sim: {cos_sim_logsumexp:.8f}"
+            )
+
+            # Performance
+            fn_cuda = lambda: light_ops.group8_int8kv_flashdecoding_diverse_stage2(
+                setup_tensors["block_seq"],
+                setup_tensors["mid_out"],
+                setup_tensors["mid_out_logsumexp"],
+                1.0 / (setup_tensors["head_dim"] ** 0.5),
+                setup_tensors["q"],
+                setup_tensors["k"],
+                setup_tensors["k_scale"],
+                setup_tensors["v"],
+                setup_tensors["v_scale"],
+                setup_tensors["Req_to_tokens"],
+                setup_tensors["B_req_idx"],
+                setup_tensors["b_seq_len"],
+                setup_tensors["b_shared_seq_len"],
+                setup_tensors["max_len_in_batch"],
+            )
+            ms_cuda = triton.testing.do_bench_cudagraph(fn_cuda, rep=100)
+
+            fn_triton = lambda: flash_decode_stage2(
+                q=setup_tensors["q"],
+                k=setup_tensors["k"],
+                k_scale=setup_tensors["k_scale"],
+                v=setup_tensors["v"],
+                v_scale=setup_tensors["v_scale"],
+                Req_to_tokens=setup_tensors["Req_to_tokens"],
+                B_req_idx=setup_tensors["B_req_idx"],
+                B_Seqlen=setup_tensors["b_seq_len"],
+                b_shared_seq_len=setup_tensors["b_shared_seq_len"],
+                max_len_in_batch=setup_tensors["max_len_in_batch"],
+                mid_out=setup_tensors["mid_out"],
+                mid_out_logsumexp=setup_tensors["mid_out_logsumexp"],
+                block_seq=setup_tensors["block_seq"],
+            )
+            ms_triton = triton.testing.do_bench_cudagraph(fn_triton, rep=100)
+
+            results.append(
+                {
+                    "batch_size": batch,
+                    "seq_len": seq,
+                    "triton_ms": ms_triton,
+                    "cuda_ms": ms_cuda,
+                }
+            )
+            print(results[-1])
+
+            del setup_tensors
+
+    print(f"\n{'='*80}")
+    print("SUMMARY - Performance Comparison")
+    print(f"{'='*80}")
+    print(f"{'batch_size':<8} {'seq_len':<12} {'triton_ms':<12} {'cuda_ms':<12} {'vs cuda':<10}")
+    print(f"{'-'*80}")
+    for r in results:
+        vs_cuda = f"{r['cuda_ms']/r['triton_ms']:.2f}x"
+        emoji = "🎉" if r["triton_ms"] < r["cuda_ms"] else ""
+        print(
+            f"{r['batch_size']:<8} {r['seq_len']:<12} {r['triton_ms']:<12.3f} {r['cuda_ms']:<12.3f}"
+            f"{vs_cuda:<10} {emoji}"
+        )
+    print(f"{'='*80}")
diff --git a/unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse_stage3.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py
similarity index 79%
rename from unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse_stage3.py
rename to unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py
index b406e2dcf5..c1a0ca1e58 100644
--- a/unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse_stage3.py
+++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py
@@ -1,6 +1,8 @@
 import pytest
 import torch
-from lightllm.models.llama.triton_kernel.ppl_int8kv_flash_decoding_diverse_stage3 import flash_diverse_decode_stage3
+from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage3 import (
+    flash_diverse_decode_stage3,
+)
 
 
 @pytest.mark.parametrize(
@@ -23,7 +25,10 @@ def test_flash_diverse_decode_stage3(batch, head_num, seq_len, shared_seq_len, b
     flash_diverse_decode_stage3(mid_out, mid_out_logexpsum, B_Seqlen, b_shared_seq_len, out, block_seq)
 
     true_out = torch.zeros_like(out)
-    from lightllm.models.llama.triton_kernel.flash_decoding_stage2 import flash_decode_stage2
+
+    from lightllm.common.basemodel.triton_kernel.att.decode_att.mha.flash_decoding.flash_decoding_stage2 import (
+        flash_decode_stage2,
+    )
 
     flash_decode_stage2(mid_out, mid_out_logexpsum, B_Seqlen, true_out, block_seq)
 
diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py
new file mode 100644
index 0000000000..d1a53f873f
--- /dev/null
+++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py
@@ -0,0 +1,104 @@
+import torch
+import time
+import pytest
+from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import (
+    context_attention_fwd_contiguous_kv,
+)
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def torch_context_attention_fwd2(q, k, v, o, b_start_loc, b_kv_start_loc, b_seq_len, b_prompt_cache_len):
+    batch = b_start_loc.shape[0]
+
+    for i in range(batch):
+        start_loc = b_start_loc[i]
+        kv_start_loc = b_kv_start_loc[i]
+        seq_len = b_seq_len[i]
+        prompt_cache_len = b_prompt_cache_len[i]
+        cur_q = q[start_loc : start_loc + seq_len - prompt_cache_len, :, :]
+        cur_q = cur_q.clone().to(torch.float32)
+        cur_k = k[kv_start_loc : (kv_start_loc + seq_len), :, :]
+        cur_k = cur_k.clone().to(torch.float32)
+
+        cur_v = v[kv_start_loc : (kv_start_loc + seq_len), :, :]
+        cur_v = cur_v.clone().to(torch.float32)
+
+        dk = cur_q.shape[-1]
+        cur_q = cur_q.permute(1, 0, 2)
+        cur_k = cur_k.permute(1, 2, 0)
+        cur_v = cur_v.permute(1, 0, 2)
+        dk = cur_q.shape[-1]
+
+        p = torch.matmul(cur_q, cur_k) / torch.sqrt(torch.tensor(dk, dtype=torch.float32))
+
+        q_index = (torch.arange(cur_q.shape[1]).to(p.device) + prompt_cache_len).view(-1, 1)
+        k_index = torch.arange(seq_len).to(p.device).view(1, -1)
+
+        p[:, (q_index < k_index)] = float("-inf")
+
+        s = torch.nn.functional.softmax(p, dim=-1)
+
+        o[start_loc : start_loc + seq_len - prompt_cache_len, :, :] = torch.matmul(s, cur_v).transpose(0, 1)
+
+
+@pytest.mark.parametrize(
+    "B, H, N_CTX, D_HEAD, prompt_cache_len",
+    [
+        (b, H, N_CTX, D_HEAD, prompt_cache_len)
+        for b in [1, 2, 4]
+        for H in [1, 8]
+        for N_CTX in [3, 10, 1024]
+        for D_HEAD in [64, 128]
+        for prompt_cache_len in [0, 56, 200]
+    ],
+)
+def test_context_attention_fwd_contiguous_kv(B, H, N_CTX, D_HEAD, prompt_cache_len):
+    dtype = torch.float16
+    prompt_cache_len = 0
+    if prompt_cache_len >= N_CTX - 1:
+        return
+
+    q = torch.empty((B * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.1, std=0.2)
+    kv = torch.empty((B * N_CTX, 2 * H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.4, std=0.2)
+    k = kv[:, :H, :]
+    v = kv[:, H:, :]
+
+    o = torch.empty((B * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda")
+    torch_o = torch.empty((B * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda")
+
+    max_q_input_len = N_CTX - prompt_cache_len
+
+    b_seq_len = torch.ones((B,), dtype=torch.int32, device="cuda")
+    b_start_loc = torch.zeros((B,), dtype=torch.int32, device="cuda")
+    b_seq_len = torch.ones((B,), dtype=torch.int32, device="cuda")
+    b_prompt_cache_len = torch.zeros(B, dtype=torch.int32, device="cuda")
+
+    for i in range(B):
+        b_seq_len[i] = N_CTX
+        if i != 0:
+            b_start_loc[i] = b_start_loc[i - 1] + N_CTX - prompt_cache_len
+        b_prompt_cache_len[i] = prompt_cache_len
+
+    b_kv_start_loc = torch.cumsum(b_seq_len, dim=0, dtype=torch.int32) - b_seq_len
+    torch_context_attention_fwd2(q, k, v, torch_o, b_start_loc, b_kv_start_loc, b_seq_len, b_prompt_cache_len)
+    context_attention_fwd_contiguous_kv(
+        q=q,
+        k=k,
+        v=v,
+        o=o,
+        b_start_loc=b_start_loc,
+        b_kv_start_loc=b_kv_start_loc,
+        b_seq_len=b_seq_len,
+        max_q_input_len=max_q_input_len,
+        b_prompt_cache_len=b_prompt_cache_len,
+    )
+
+    assert torch.allclose(torch_o, o, atol=1e-2, rtol=0)
+    cos = torch.nn.CosineSimilarity(0)
+    assert cos(o.flatten().float(), torch_o.flatten().float()) > 0.99
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/unit_tests/models/llama/test_context_flashattention_nopad.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py
similarity index 92%
rename from unit_tests/models/llama/test_context_flashattention_nopad.py
rename to unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py
index f24ab619bd..541594306d 100644
--- a/unit_tests/models/llama/test_context_flashattention_nopad.py
+++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py
@@ -5,12 +5,11 @@
 import torch.nn.functional as F
 import flashinfer
 from lightllm.utils.log_utils import init_logger
-from lightllm.models.llama.triton_kernel.context_flashattention_nopad import (
+from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import (
     context_attention_fwd,
     context_attention_fwd_no_prompt_cache,
 )
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.common.req_manager import ReqManager
 
 logger = init_logger(__name__)
 
@@ -54,14 +53,14 @@ def test_context_attention_fwd(batch, seqlen, q_heads, kv_heads, head_dim):
 
     infer_state = LlamaInferStateInfo()
     infer_state.batch_size = Z
-    infer_state.max_len_in_batch = N_CTX
+    infer_state.max_q_seq_len = N_CTX
     infer_state.total_token_num = Z * N_CTX
-    infer_state.req_manager = ReqManager(Z, N_CTX, None)
+    infer_state.req_manager = type("Object", (), {})()
     infer_state.req_manager.req_to_token_indexs = req_to_token_indexs
     infer_state.b_req_idx = b_req_idx
     infer_state.b_seq_len = b_seq_len
     infer_state.b_ready_cache_len = b_ready_cache_len
-    infer_state.b_start_loc = q_start_loc
+    infer_state.b_q_start_loc = q_start_loc
 
     context_attention_fwd(
         q,
@@ -69,10 +68,10 @@ def test_context_attention_fwd(batch, seqlen, q_heads, kv_heads, head_dim):
         kv[:, KV_HEADS:, :],
         o,
         infer_state.b_req_idx,
-        infer_state.b_start_loc,
+        infer_state.b_q_start_loc,
         infer_state.b_seq_len,
         infer_state.b_ready_cache_len,
-        infer_state.max_len_in_batch,
+        infer_state.max_q_seq_len,
         infer_state.req_manager.req_to_token_indexs,
     )
 
@@ -127,7 +126,11 @@ def test_context_attention_fwd(batch, seqlen, q_heads, kv_heads, head_dim):
     "batch, seqlen, q_heads, kv_heads, head_dim",
     [
         (a, b, c, d, e)
-        for a in [1, 16, 32, 128, 512]
+        for a in [
+            1,
+            16,
+            32,
+        ]
         for b in [16, 32, 512, 1024]
         for c in [28]
         for d in [4]
@@ -149,18 +152,18 @@ def test_context_attention_fwd_no_prompt_cache(batch, seqlen, q_heads, kv_heads,
 
     infer_state = LlamaInferStateInfo()
     infer_state.batch_size = Z
-    infer_state.max_len_in_batch = N_CTX
+    infer_state.max_q_seq_len = N_CTX
     infer_state.b_seq_len = b_seq_len
-    infer_state.b_start_loc = b_start_loc
+    infer_state.b_q_start_loc = b_start_loc
 
     context_attention_fwd_no_prompt_cache(
         q,
         k,
         v,
         o,
-        infer_state.b_start_loc,
+        infer_state.b_q_start_loc,
         infer_state.b_seq_len,
-        infer_state.max_len_in_batch,
+        infer_state.max_q_seq_len,
     )
 
     head_dim = HEAD_DIM
diff --git a/unit_tests/models/deepseek2/test_destindex_copy_kv.py b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_mla_destindex_copy_kv.py
similarity index 93%
rename from unit_tests/models/deepseek2/test_destindex_copy_kv.py
rename to unit_tests/common/basemodel/triton_kernel/kv_copy/test_mla_destindex_copy_kv.py
index 1379dc72de..ed0c6e369f 100644
--- a/unit_tests/models/deepseek2/test_destindex_copy_kv.py
+++ b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_mla_destindex_copy_kv.py
@@ -1,6 +1,6 @@
 import torch
 import pytest
-from lightllm.models.deepseek2.triton_kernel.destindex_copy_kv import destindex_copy_kv
+from lightllm.common.basemodel.triton_kernel.kv_copy.mla_copy_kv import destindex_copy_kv
 from lightllm.utils.log_utils import init_logger
 import torch.nn.functional as F
 
diff --git a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int4kv_copy_kv.py b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int4kv_copy_kv.py
new file mode 100644
index 0000000000..83537ec708
--- /dev/null
+++ b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int4kv_copy_kv.py
@@ -0,0 +1,62 @@
+import torch
+import pytest
+import numpy as np
+from typing import Tuple
+from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import destindex_copy_int4kv, dequantize_int4kv
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def test_quanted_and_dequant():
+    """Test quantization followed by dequantization."""
+    batch_size = 1
+    seq_len = 8
+    head_num = 4
+    k_head_num = 2
+    v_head_num = 2
+    assert k_head_num + v_head_num == head_num
+    head_dim = 64
+    quant_group_size = 8
+
+    # Create original data
+    original_kv = torch.randn(batch_size * seq_len, head_num, head_dim, dtype=torch.float32).clamp_(-1, 1).cuda()
+    dest_loc = torch.arange(batch_size * seq_len, dtype=torch.int64).cuda()
+
+    # Quantize
+    group_count = head_dim // quant_group_size
+    kv_buffer = torch.zeros(batch_size * seq_len, head_num, head_dim // 2, dtype=torch.int8).cuda()
+    kv_scale_buffer = torch.zeros(batch_size * seq_len, head_num, group_count, dtype=torch.float32).cuda()
+    destindex_copy_int4kv(original_kv, dest_loc, kv_buffer, kv_scale_buffer, quant_group_size)
+
+    # Dequantize
+    req_to_token_indexs = torch.arange(seq_len, dtype=torch.int64).view(1, -1).cuda()
+    b_seq_len = torch.tensor([seq_len], dtype=torch.int32).cuda()
+    b_req_idx = torch.tensor([0], dtype=torch.int32).cuda()
+    b_kv_start_loc = torch.tensor([0], dtype=torch.int32).cuda()
+
+    recovered_kv = torch.zeros(batch_size * seq_len, head_num, head_dim, dtype=torch.float32).cuda()
+
+    dequantize_int4kv(
+        k=kv_buffer[:, 0:k_head_num, :],
+        k_scale=kv_scale_buffer[:, 0:k_head_num, :],
+        v=kv_buffer[:, k_head_num:, :],
+        v_scale=kv_scale_buffer[:, k_head_num:, :],
+        req_to_token_indexs=req_to_token_indexs,
+        b_seq_len=b_seq_len,
+        b_req_idx=b_req_idx,
+        b_kv_start_loc=b_kv_start_loc,
+        k_out=recovered_kv[:, :k_head_num, :],
+        v_out=recovered_kv[:, k_head_num:, :],
+        max_len_in_batch=seq_len,
+        quant_group_size=quant_group_size,
+    )
+
+    logger.info("Round-trip test completed!")
+    assert torch.allclose(recovered_kv, original_kv, atol=2 / 14, rtol=0)
+    cos = torch.nn.CosineSimilarity(0)
+    assert cos(recovered_kv.flatten().float(), original_kv.flatten().float()) > 0.99
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int8kv_copy_kv.py b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int8kv_copy_kv.py
new file mode 100644
index 0000000000..149c9894ae
--- /dev/null
+++ b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int8kv_copy_kv.py
@@ -0,0 +1,86 @@
+import torch
+import time
+import pytest
+from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import (
+    dequantize_int8kv,
+    destindex_copy_quantize_kv,
+)
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def torch_dequant(kv, kv_scale, b_req_idx, b_seq_len, req_to_token_indexs, odtype, group_quant_size):
+    batch = b_req_idx.shape[0]
+    tmp_out = []
+    for i in range(batch):
+        req_idx = b_req_idx[i]
+        seq_len = b_seq_len[i]
+        kv_loc = req_to_token_indexs[req_idx, :seq_len]
+        head_num = kv.shape[1]
+        cur_kv = kv[kv_loc, :, :].reshape(seq_len, head_num, -1, group_quant_size).to(odtype)
+        cur_scale = kv_scale[kv_loc, :, :].reshape(seq_len, head_num, -1, 1)
+        out = cur_kv * cur_scale
+        tmp_out.append(out.reshape(seq_len, head_num, -1))
+    return torch.cat(tmp_out, dim=0)
+
+
+@pytest.mark.parametrize(
+    "B, H, N_CTX, D_HEAD, group_quant_size",
+    [
+        (b, H, N_CTX, D_HEAD, group_quant_size)
+        for b in [1, 2, 4]
+        for H in [1, 8]
+        for N_CTX in [3, 10, 1024]
+        for D_HEAD in [64, 128]
+        for group_quant_size in [8, 16]
+    ],
+)
+def test_dequantize_int8kv(B, H, N_CTX, D_HEAD, group_quant_size):
+    dtype = torch.bfloat16
+    kv = torch.empty((B * N_CTX, 2 * H, D_HEAD), dtype=torch.int8, device="cuda").random_(-10, 10)
+    kv_scale = torch.randn((B * N_CTX, 2 * H, D_HEAD // group_quant_size), dtype=dtype, device="cuda")
+    out = torch.empty((B * N_CTX, 2 * H, D_HEAD), dtype=dtype, device="cuda")
+    req_to_token_indexs = torch.empty((B, N_CTX), dtype=torch.int32, device="cuda")
+    max_input_len = N_CTX
+    b_seq_len = torch.ones((B,), dtype=torch.int32, device="cuda")
+    b_seq_len.fill_(N_CTX)
+    b_req_idx = torch.arange(0, B, dtype=torch.int32, device="cuda")
+    req_to_token_indexs.view(-1)[:] = torch.arange(0, B * N_CTX, dtype=torch.int32, device="cuda")
+    b_kv_start_loc = torch.cumsum(b_seq_len, dim=0, dtype=torch.int32) - b_seq_len
+
+    k = kv[:, :H, :]
+    v = kv[:, H:, :]
+    k_scale = kv_scale[:, :H, :]
+    v_scale = kv_scale[:, H:, :]
+
+    ground_out = torch_dequant(
+        kv=kv,
+        kv_scale=kv_scale,
+        b_req_idx=b_req_idx,
+        b_seq_len=b_seq_len,
+        req_to_token_indexs=req_to_token_indexs,
+        odtype=out.dtype,
+        group_quant_size=group_quant_size,
+    )
+    dequantize_int8kv(
+        k=k,
+        k_scale=k_scale,
+        v=v,
+        v_scale=v_scale,
+        req_to_token_indexs=req_to_token_indexs,
+        b_seq_len=b_seq_len,
+        b_req_idx=b_req_idx,
+        b_kv_start_loc=b_kv_start_loc,
+        k_out=out[:, :H, :],
+        v_out=out[:, H:, :],
+        max_len_in_batch=max_input_len,
+        quant_group_size=group_quant_size,
+    )
+    assert torch.allclose(out, ground_out, atol=1e-2, rtol=0)
+    cos = torch.nn.CosineSimilarity(0)
+    assert cos(out.flatten().float(), ground_out.flatten().float()) > 0.99
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/unit_tests/models/deepseek2/test_gqa_flash_decoding.py b/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py
similarity index 92%
rename from unit_tests/models/deepseek2/test_gqa_flash_decoding.py
rename to unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py
index d0bc670ecb..a5ac9708d2 100644
--- a/unit_tests/models/deepseek2/test_gqa_flash_decoding.py
+++ b/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py
@@ -5,9 +5,10 @@
 import torch.nn.functional as F
 import flashinfer
 from lightllm.utils.log_utils import init_logger
-from lightllm.models.deepseek2.triton_kernel.gqa_flash_decoding import gqa_token_decode_attention_flash_decoding
+from lightllm.common.basemodel.triton_kernel.mla_att.decode_att.gqa_flash_decoding import (
+    gqa_token_decode_attention_flash_decoding,
+)
 from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
-from lightllm.common.req_manager import ReqManager
 
 logger = init_logger(__name__)
 
@@ -53,7 +54,7 @@ def test_gqa_flash_decoding(batch, seqlen, heads, nope_head, rope_head):
     infer_state.batch_size = Z
     infer_state.max_len_in_batch = N_CTX
     infer_state.total_token_num = Z * N_CTX
-    infer_state.req_manager = ReqManager(Z, N_CTX, None)
+    infer_state.req_manager = type("Object", (), {})()
     infer_state.req_manager.req_to_token_indexs = req_to_token_indexs
     infer_state.b_req_idx = b_req_idx
     infer_state.b_seq_len = b_seq_len
@@ -67,10 +68,6 @@ def test_gqa_flash_decoding(batch, seqlen, heads, nope_head, rope_head):
         kv_nope,
         kv_rope,
         infer_state,
-        H,
-        D_HEAD,
-        ROPE_HEAD,
-        D_HEAD,
         sm_scale,
         o,
     )
diff --git a/unit_tests/common/basemodel/triton_kernel/test_atomic_event.py b/unit_tests/common/basemodel/triton_kernel/test_atomic_event.py
index 536cad90fc..0afcd5558a 100644
--- a/unit_tests/common/basemodel/triton_kernel/test_atomic_event.py
+++ b/unit_tests/common/basemodel/triton_kernel/test_atomic_event.py
@@ -18,10 +18,10 @@ def test_add_in_place():
     assert input.item() == 3, "最终值应为 3"
 
 
-@pytest.mark.timeout(2)
-def test_wait_timeout():
-    input = torch.zeros((1,), device="cuda", dtype=torch.int32)
-    wait_value(input, 4)
+# @pytest.mark.timeout(2)
+# def test_wait_timeout():
+#     input = torch.zeros((1,), device="cuda", dtype=torch.int32)
+#     wait_value(input, 4)
 
 
 if __name__ == "__main__":
diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py b/unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py
index 5c3ca89c65..41bc217b94 100644
--- a/unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py
+++ b/unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py
@@ -1,21 +1,9 @@
 import torch
 import pytest
-import easydict
 from lightllm.common.basemodel.triton_kernel.gen_decode_params import gen_decode_params
-from lightllm.utils.envs_utils import set_env_start_args
 
 
 def test_gen_decode_params_basic():
-    set_env_start_args(
-        easydict.EasyDict(
-            {
-                "mtp_step": 0,
-                "enable_flashinfer_prefill": False,
-                "enable_flashinfer_decode": False,
-            }
-        )
-    )
-
     b_seq_len = torch.ones((9,), dtype=torch.int64, device="cuda") * 8192
     (
         b_q_seq_len,
diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_sampling_params.py b/unit_tests/common/basemodel/triton_kernel/test_gen_sampling_params.py
index 99971dea25..e9d0193279 100644
--- a/unit_tests/common/basemodel/triton_kernel/test_gen_sampling_params.py
+++ b/unit_tests/common/basemodel/triton_kernel/test_gen_sampling_params.py
@@ -25,6 +25,7 @@ def test_token_id_counter():
     for _ in range(100):
         token_id_counter(prompt_ids=test_prompt_ids, out_token_id_counter=test_token_id_counter)
     end_event.record()
+    end_event.synchronize()
     logger.info(f"test_token_id_count cost time: {start_event.elapsed_time(end_event)} ms")
 
 
diff --git a/unit_tests/models/deepseek2/test_repack_kv_index.py b/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py
similarity index 96%
rename from unit_tests/models/deepseek2/test_repack_kv_index.py
rename to unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py
index f9e5928a9e..b5184d3caa 100644
--- a/unit_tests/models/deepseek2/test_repack_kv_index.py
+++ b/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py
@@ -1,7 +1,7 @@
 import torch
 import pytest
 from lightllm.utils.log_utils import init_logger
-from lightllm.models.deepseek2.triton_kernel.repack_kv_index import repack_kv_index
+from lightllm.common.basemodel.triton_kernel.repack_kv_index import repack_kv_index
 
 logger = init_logger(__name__)
 
diff --git a/unit_tests/common/fused_moe/test_deepep.py b/unit_tests/common/fused_moe/test_deepep.py
index c846be0961..45778244b7 100644
--- a/unit_tests/common/fused_moe/test_deepep.py
+++ b/unit_tests/common/fused_moe/test_deepep.py
@@ -1,12 +1,13 @@
+import pytest
+
+pytest.skip(reason="need special env, install deep_ep and deep_gemm", allow_module_level=True)
+
 import os
 import torch
 import torch.distributed as dist
-import pytest
 import deep_ep
 import random
 import numpy as np
-from deep_ep import Buffer, EventOverlap
-from deep_gemm import ceil_div, get_col_major_tma_aligned_tensor
 from lightllm.common.fused_moe.grouped_fused_moe_ep import fused_experts_impl
 from lightllm.common.fused_moe.deepep_scatter_gather import ep_scatter, ep_gather
 from typing import Tuple
@@ -25,6 +26,8 @@
 def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     assert x.dim() == 2
     m, n = x.shape
+    from deep_gemm import ceil_div
+
     x_padded = torch.zeros((ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device)
     x_padded[:m, :n] = x
     x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
diff --git a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py
index eba15b2a1c..ab2cc4976a 100644
--- a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py
+++ b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py
@@ -1,10 +1,22 @@
 import torch
-import time
 import pytest
+
+
+def is_fp8_native_supported():
+    """检查是否为 H100/B200 等原生支持 FP8 的硬件 (SM90+)"""
+    if not torch.cuda.is_available():
+        return False
+    major, _ = torch.cuda.get_device_capability()
+    return major >= 9
+
+
+if not is_fp8_native_supported():
+    pytest.skip(reason="not support fp8 test in this gpu card", allow_module_level=True)
+
 import random
 from lightllm.common.fused_moe.moe_silu_and_mul_mix_quant_ep import silu_and_mul_masked_post_quant_fwd
 from lightllm.common.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
-from lightllm.common.quantization.triton_quant.fp8.fp8act_quant_kernel import per_token_group_quant_fp8
+from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
diff --git a/unit_tests/common/fused_moe/test_softmax_topk.py b/unit_tests/common/fused_moe/test_softmax_topk.py
index 262c37a0f2..6252dfa8c3 100755
--- a/unit_tests/common/fused_moe/test_softmax_topk.py
+++ b/unit_tests/common/fused_moe/test_softmax_topk.py
@@ -9,7 +9,10 @@
 
 
 def benchmark(M, N, K, renorm, runs):
-    import sgl_kernel as sgl_ops
+    try:
+        import sgl_kernel as sgl_ops
+    except Exception as e:
+        pytest.skip(f"no sgl_kernel error: {str(e)}", allow_module_level=True)
 
     gating = torch.randn(M, N, device="cuda", dtype=torch.float32)
     torch.cuda.synchronize()
diff --git a/unit_tests/common/quantization/test_fp8_scaled_mm_per_token.py b/unit_tests/common/quantization/test_fp8_scaled_mm_per_token.py
index 1ddb20b632..e6a0d52c72 100644
--- a/unit_tests/common/quantization/test_fp8_scaled_mm_per_token.py
+++ b/unit_tests/common/quantization/test_fp8_scaled_mm_per_token.py
@@ -1,7 +1,19 @@
 import torch
 import pytest
 import torch.nn.functional as F
-from lightllm.common.quantization.triton_quant.fp8.fp8w8a8_scaled_mm_per_token_kernel import fp8_scaled_mm_per_token
+from lightllm.common.basemodel.triton_kernel.quantization.scaled_mm_per_token_kernel import fp8_scaled_mm_per_token
+
+
+def is_fp8_native_supported():
+    """检查是否为 H100/B200 等原生支持 FP8 的硬件 (SM90+)"""
+    if not torch.cuda.is_available():
+        return False
+    major, _ = torch.cuda.get_device_capability()
+    return major >= 9
+
+
+if not is_fp8_native_supported():
+    pytest.skip("not support fp8 in this gpu card", allow_module_level=True)
 
 
 @pytest.mark.parametrize("M", [1, 2, 4, 8, 16, 32, 64, 128])
diff --git a/unit_tests/models/deepseek2/test_destindex_copy_kv_fp8.py b/unit_tests/models/deepseek2/test_destindex_copy_kv_fp8.py
deleted file mode 100644
index 4f9c0a3373..0000000000
--- a/unit_tests/models/deepseek2/test_destindex_copy_kv_fp8.py
+++ /dev/null
@@ -1,48 +0,0 @@
-import torch
-import pytest
-from lightllm.models.deepseek2.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8
-from lightllm.utils.log_utils import init_logger
-import torch.nn.functional as F
-
-logger = init_logger(__name__)
-
-seed = 42
-torch.manual_seed(seed)
-
-if torch.cuda.is_available():
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-
-@pytest.mark.parametrize(
-    "batch, seqlen, heads, nope_head, rope_head, copy_len",
-    [
-        (a, b, c, d, e, f)
-        for a in [1, 16, 32, 128, 512]
-        for b in [1024, 2048]
-        for c in [1]
-        for d in [512]
-        for e in [64]
-        for f in [10, 20, 100, 1024]
-    ],
-)
-def test_destindex_copy_kv_fp8(batch, seqlen, heads, nope_head, rope_head, copy_len):
-    B, N_CTX, H, NOPE_HEAD, ROPE_HEAD, COPY_LEN = batch, seqlen, heads, nope_head, rope_head, copy_len
-    dtype = torch.bfloat16
-    NUM = COPY_LEN
-    dest_loc = torch.arange(NUM).cuda()
-    kv = torch.randn((len(dest_loc), H, NOPE_HEAD + ROPE_HEAD), dtype=dtype).cuda()
-    out = torch.zeros((B * N_CTX, H, NOPE_HEAD + ROPE_HEAD + 2), dtype=torch.uint8).cuda()
-
-    fp8_type = torch.float8_e4m3fn
-    kv_nope = kv[:, :, :NOPE_HEAD]
-    kv_rope = kv[:, :, NOPE_HEAD:]
-    O_nope = out[:, :, :NOPE_HEAD].view(fp8_type)
-    O_rope = out[:, :, NOPE_HEAD:-2].view(fp8_type)
-    O_scale = out[:, :, -2:].view(dtype)
-    destindex_copy_kv_fp8(kv_nope, kv_rope, dest_loc, O_nope, O_rope, O_scale)
-
-    cos1 = F.cosine_similarity(O_nope[:NUM].to(dtype) * O_scale[:NUM], kv_nope).mean()
-    cos2 = F.cosine_similarity(O_rope[:NUM].to(dtype) * O_scale[:NUM], kv_rope).mean()
-    assert cos1 > 0.98
-    assert cos2 > 0.98
diff --git a/unit_tests/models/deepseek2/test_gqa_flash_decoding_fp8.py b/unit_tests/models/deepseek2/test_gqa_flash_decoding_fp8.py
deleted file mode 100644
index 72d9d9accc..0000000000
--- a/unit_tests/models/deepseek2/test_gqa_flash_decoding_fp8.py
+++ /dev/null
@@ -1,79 +0,0 @@
-import torch
-import pytest
-import numpy as np
-import torch.nn.functional as F
-from lightllm.utils.log_utils import init_logger
-from lightllm.models.deepseek2.triton_kernel.gqa_flash_decoding import gqa_token_decode_attention_flash_decoding
-from lightllm.models.deepseek2.triton_kernel.gqa_flash_decoding_fp8 import gqa_token_decode_attention_flash_decoding_fp8
-from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
-from lightllm.common.req_manager import ReqManager
-
-logger = init_logger(__name__)
-
-seed = 42
-torch.manual_seed(seed)
-
-if torch.cuda.is_available():
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-
-@pytest.mark.parametrize(
-    "batch, seqlen, heads, nope_head, rope_head",
-    [(a, b, c, d, e) for a in [1, 16, 32, 128] for b in [16, 32, 512, 2048] for c in [16] for d in [512] for e in [64]],
-)
-def test_gqa_flash_decoding_fp8(batch, seqlen, heads, nope_head, rope_head):
-    Z, N_CTX, H, D_HEAD, ROPE_HEAD = batch, seqlen, heads, nope_head, rope_head
-    dtype = torch.bfloat16
-    sm_scale = 1.0 / ((D_HEAD + ROPE_HEAD) ** 0.5)
-    q = torch.randn((Z, H, D_HEAD), dtype=dtype, device="cuda")
-    q_rope = torch.randn((Z, H, ROPE_HEAD), dtype=dtype, device="cuda")
-
-    kv = torch.randn((Z * N_CTX, 1, D_HEAD + ROPE_HEAD), dtype=dtype, device="cuda")
-    kv_scale = torch.randn((Z * N_CTX, 1, 1), dtype=dtype, device="cuda")
-    kv_fp8 = kv.to(torch.float8_e4m3fn)
-
-    req_to_token_indexs = torch.zeros((10, Z * N_CTX), dtype=torch.int32, device="cuda")
-    b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda")
-    b_req_idx = torch.ones((Z,), dtype=torch.int32, device="cuda")
-
-    b_seq_len[0] = N_CTX
-    b_req_idx[0] = 0
-    req_to_token_indexs[0][:N_CTX] = torch.tensor(np.arange(N_CTX), dtype=torch.int32).cuda()
-
-    o = torch.empty((Z, H, D_HEAD), dtype=dtype, device="cuda")
-    o1 = torch.empty((Z, H, D_HEAD), dtype=dtype, device="cuda")
-
-    infer_state = Deepseek2InferStateInfo()
-    infer_state.batch_size = Z
-    infer_state.max_len_in_batch = N_CTX
-    infer_state.total_token_num = Z * N_CTX
-    infer_state.req_manager = ReqManager(Z, N_CTX, None)
-    infer_state.req_manager.req_to_token_indexs = req_to_token_indexs
-    infer_state.b_req_idx = b_req_idx
-    infer_state.b_seq_len = b_seq_len
-
-    kv_nope = kv_fp8[:, :, :D_HEAD].to(dtype) * kv_scale
-    kv_rope = kv_fp8[:, :, D_HEAD:].to(dtype) * kv_scale
-    gqa_token_decode_attention_flash_decoding(
-        q,
-        q_rope,
-        kv_nope,
-        kv_rope,
-        infer_state,
-        H,
-        D_HEAD,
-        ROPE_HEAD,
-        D_HEAD,
-        sm_scale,
-        o,
-    )
-
-    kv_nope_fp8 = kv_fp8[:, :, :D_HEAD]
-    kv_rope_fp8 = kv_fp8[:, :, D_HEAD:]
-    gqa_token_decode_attention_flash_decoding_fp8(
-        q, q_rope, kv_nope_fp8, kv_rope_fp8, kv_scale, infer_state, H, D_HEAD, ROPE_HEAD, D_HEAD, sm_scale, o1
-    )
-
-    cos_sim = F.cosine_similarity(o, o1).mean()
-    assert cos_sim > 0.99
diff --git a/unit_tests/models/llama/test_context_flashattention_nopad_fa3_fp8.py b/unit_tests/models/llama/test_context_flashattention_nopad_fa3_fp8.py
deleted file mode 100644
index 737bb655b1..0000000000
--- a/unit_tests/models/llama/test_context_flashattention_nopad_fa3_fp8.py
+++ /dev/null
@@ -1,154 +0,0 @@
-import torch
-import time
-import pytest
-import triton as tl
-import numpy as np
-import torch.nn.functional as F
-from lightllm.utils.log_utils import init_logger
-from lightllm.models.llama.triton_kernel.context_flashattention_nopad import (
-    context_attention_fwd,
-)
-from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.utils.sgl_utils import flash_attn_with_kvcache
-from lightllm.common.basemodel.triton_kernel.q_per_head_fp8_quant import q_per_head_fp8_quant
-
-logger = init_logger(__name__)
-
-seed = 42
-torch.manual_seed(seed)
-
-if torch.cuda.is_available():
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-
-def kv_quantize_per_head_fp8(kv_buffer: torch.Tensor, seq_lens):
-    device = kv_buffer.device
-    B = seq_lens.size(0)
-    min_fp8 = torch.finfo(torch.float8_e4m3fn).min
-    max_fp8 = torch.finfo(torch.float8_e4m3fn).max
-    _, S_max, H, D = kv_buffer.shape
-    seq_range = torch.arange(S_max, device=device)[None, :]
-    valid_mask = (seq_range < seq_lens[:, None]).view(B, S_max, 1, 1)
-    masked = kv_buffer * valid_mask
-    max_per_bh = masked.abs().amax(dim=(1, 3))  # [B, H]
-    scales = torch.where(max_per_bh > 0, max_per_bh / max_fp8, torch.ones_like(max_per_bh)).to(torch.float32)
-    scales_exp = scales.view(B, 1, H, 1)
-    q = (kv_buffer / scales_exp).clamp(min_fp8, max_fp8).to(torch.float8_e4m3fn)
-    return q, scales
-
-
-@pytest.mark.parametrize(
-    "batch, seqlen, q_heads, kv_heads, head_dim",
-    [
-        (a, b, c, d, e)
-        for a in [1, 16, 32, 128, 512]
-        for b in [16, 32, 512, 1024]
-        for c in [28]
-        for d in [4]
-        for e in [128]
-    ],
-)
-def test_context_attention_fwd_fa3_fp8(batch, seqlen, q_heads, kv_heads, head_dim):
-    Z, N_CTX, Q_HEADS, KV_HEADS, HEAD_DIM = batch, seqlen, q_heads, kv_heads, head_dim
-    dtype = torch.bfloat16
-    kv = torch.randn((Z * N_CTX, 2 * KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    # for i in range(Z * N_CTX):
-    #     kv[i] = torch.randn((2 * KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda") * (i % 10 + 1)
-
-    max_input_len = Z * N_CTX
-    req_to_token_indexs = torch.randperm(max_input_len, dtype=torch.int32).cuda().view(Z, N_CTX)
-    b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * (N_CTX // 2)
-    rand_num = torch.randint_like(b_seq_len, high=(N_CTX // 2), dtype=torch.int32, device="cuda")
-    b_seq_len += rand_num
-    b_ready_cache_len = torch.zeros_like(b_seq_len, dtype=torch.int32, device="cuda")
-    if N_CTX > 1:
-        b_ready_cache_len = torch.randint_like(b_seq_len, high=(N_CTX - 1) // 2, dtype=torch.int32, device="cuda")
-    b_req_idx = torch.randperm(Z, dtype=torch.int32).cuda()
-    q_lens = b_seq_len - b_ready_cache_len
-    q_start_loc = q_lens.cumsum(0) - q_lens
-
-    q = torch.randn((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    o = torch.zeros((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    o1 = torch.zeros((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-
-    infer_state = LlamaInferStateInfo()
-    infer_state.batch_size = Z
-    infer_state.max_len_in_batch = N_CTX
-    infer_state.total_token_num = Z * N_CTX
-    infer_state.b_req_idx = b_req_idx
-    infer_state.b_seq_len = b_seq_len
-    infer_state.b_ready_cache_len = b_ready_cache_len
-    infer_state.b_start_loc = q_start_loc
-
-    context_attention_fwd(
-        q,
-        kv[:, :KV_HEADS, :],
-        kv[:, KV_HEADS:, :],
-        o,
-        infer_state.b_req_idx,
-        infer_state.b_start_loc,
-        infer_state.b_seq_len,
-        infer_state.b_ready_cache_len,
-        infer_state.max_len_in_batch,
-        req_to_token_indexs,
-    )
-
-    batch_size = Z
-    head_dim = HEAD_DIM
-    q_heads = Q_HEADS
-    kv_heads = KV_HEADS
-    page_table = torch.empty((batch_size, N_CTX), dtype=torch.int32, device="cuda")
-    page_table.copy_(req_to_token_indexs[b_req_idx, :N_CTX])
-
-    q_starts = torch.zeros((Z + 1,)).int().cuda()
-    q_starts[1:] = torch.cumsum(b_seq_len - b_ready_cache_len, dim=0)
-    kv_starts = torch.zeros_like(q_starts)
-    kv_starts[1:] = torch.cumsum(b_seq_len, dim=0)
-
-    k_cache = kv[:, :KV_HEADS, :]
-    v_cache = kv[:, KV_HEADS:, :]
-    # o1 = flash_attn_with_kvcache(
-    #     q=q,
-    #     k_cache=k_cache.reshape(-1, 1, kv_heads, head_dim),
-    #     v_cache=v_cache.reshape(-1, 1, kv_heads, head_dim),
-    #     page_table=page_table,
-    #     cache_seqlens=infer_state.b_seq_len,
-    #     cu_seqlens_q=q_starts,
-    #     cu_seqlens_k_new=kv_starts,
-    #     max_seqlen_q=N_CTX,
-    #     causal=True,
-    #     window_size=(-1, -1),
-    #     softcap=0.0,
-    #     return_softmax_lse=False,
-    # )
-
-    q, q_scale = q_per_head_fp8_quant(q.view(q.shape[0], kv_heads, -1), q_lens, q_starts)
-    k, k_scale = kv_quantize_per_head_fp8(k_cache[page_table], b_seq_len)
-    v, v_scale = kv_quantize_per_head_fp8(v_cache[page_table], b_seq_len)
-    o1 = flash_attn_with_kvcache(
-        q=q.view(-1, q_heads, head_dim),
-        k_cache=k.view(-1, N_CTX, kv_heads, head_dim).to(torch.float8_e4m3fn),
-        v_cache=v.view(-1, N_CTX, kv_heads, head_dim).to(torch.float8_e4m3fn),
-        # page_table=page_table,
-        cache_seqlens=infer_state.b_seq_len,
-        cu_seqlens_q=q_starts,
-        cu_seqlens_k_new=kv_starts,
-        max_seqlen_q=N_CTX,
-        causal=True,
-        window_size=(-1, -1),
-        softcap=0.0,
-        q_descale=q_scale.view(batch_size, kv_heads),
-        k_descale=k_scale.view(batch_size, kv_heads),
-        v_descale=v_scale.view(batch_size, kv_heads),
-        return_softmax_lse=False,
-    )
-
-    # assert torch.allclose(o, o1, atol=1e-1, rtol=1e-1)
-    cos_sim1 = F.cosine_similarity(o, o1).mean()
-    print(cos_sim1)
-    assert cos_sim1.item() == 1
-
-
-if __name__ == "__main__":
-    test_context_attention_fwd_fa3_fp8(32, 16384, 32, 4, 128)
diff --git a/unit_tests/models/llama/test_context_flashattention_nopad_flashinfer_fp8.py b/unit_tests/models/llama/test_context_flashattention_nopad_flashinfer_fp8.py
deleted file mode 100644
index 5ee2306adf..0000000000
--- a/unit_tests/models/llama/test_context_flashattention_nopad_flashinfer_fp8.py
+++ /dev/null
@@ -1,145 +0,0 @@
-import torch
-import time
-import pytest
-import numpy as np
-import torch.nn.functional as F
-import flashinfer
-from lightllm.utils.log_utils import init_logger
-from lightllm.models.llama.triton_kernel.context_flashattention_nopad import (
-    context_attention_fwd,
-)
-from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops
-
-if HAS_VLLM:
-    scaled_fp8_quant = vllm_ops.scaled_fp8_quant
-else:
-    scaled_fp8_quant = None
-
-logger = init_logger(__name__)
-
-seed = 42
-torch.manual_seed(seed)
-
-if torch.cuda.is_available():
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-
-@pytest.mark.parametrize(
-    "batch, seqlen, q_heads, kv_heads, head_dim",
-    [
-        (a, b, c, d, e)
-        for a in [1, 16, 32, 128, 512]
-        for b in [16, 32, 512, 1024]
-        for c in [28]
-        for d in [4]
-        for e in [128]
-    ],
-)
-def test_context_attention_fwd_flashinfer_fp8(batch, seqlen, q_heads, kv_heads, head_dim):
-    Z, N_CTX, Q_HEADS, KV_HEADS, HEAD_DIM = batch, seqlen, q_heads, kv_heads, head_dim
-    dtype = torch.bfloat16
-    kv = torch.randn((Z * N_CTX, 2 * KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    # for i in range(Z * N_CTX):
-    #     kv[i] = torch.randn((2 * KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda") * (i % 64 + 1)
-
-    max_input_len = Z * N_CTX
-    req_to_token_indexs = torch.randperm(max_input_len, dtype=torch.int32).cuda().view(Z, N_CTX)
-    b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * (N_CTX // 2)
-    rand_num = torch.randint_like(b_seq_len, high=(N_CTX // 2), dtype=torch.int32, device="cuda")
-    b_seq_len += rand_num
-    b_ready_cache_len = torch.zeros_like(b_seq_len, dtype=torch.int32, device="cuda")
-    if N_CTX > 1:
-        b_ready_cache_len = torch.randint_like(b_seq_len, high=(N_CTX - 1) // 2, dtype=torch.int32, device="cuda")
-    b_req_idx = torch.randperm(Z, dtype=torch.int32).cuda()
-    q_lens = b_seq_len - b_ready_cache_len
-    q_start_loc = q_lens.cumsum(0) - q_lens
-    kv_start_loc = b_seq_len.cumsum(0) - b_seq_len
-
-    q = torch.randn((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    o = torch.zeros((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    o1 = torch.zeros((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-
-    infer_state = LlamaInferStateInfo()
-    infer_state.batch_size = Z
-    infer_state.max_len_in_batch = N_CTX
-    infer_state.total_token_num = Z * N_CTX
-    infer_state.b_req_idx = b_req_idx
-    infer_state.b_seq_len = b_seq_len
-    infer_state.b_ready_cache_len = b_ready_cache_len
-    infer_state.b_start_loc = q_start_loc
-
-    context_attention_fwd(
-        q,
-        kv[:, :KV_HEADS, :],
-        kv[:, KV_HEADS:, :],
-        o,
-        infer_state.b_req_idx,
-        infer_state.b_start_loc,
-        infer_state.b_seq_len,
-        infer_state.b_ready_cache_len,
-        infer_state.max_len_in_batch,
-        req_to_token_indexs,
-    )
-
-    batch_size = Z
-    head_dim = HEAD_DIM
-    q_heads = Q_HEADS
-    kv_heads = KV_HEADS
-    page_size = 1
-    workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8).to(0)
-    q_starts = torch.zeros((Z + 1,)).int().cuda()
-    q_starts[1:] = torch.cumsum(b_seq_len - b_ready_cache_len, dim=0)
-    kv_starts = torch.zeros_like(q_starts)
-    kv_starts[1:] = torch.cumsum(b_seq_len, dim=0)
-    q_indptr = q_starts.int()
-    kv_indptr = kv_starts.int()
-    kv_indices = torch.arange(Z * N_CTX).cuda().int()
-    for b, sl, start in zip(b_req_idx, b_seq_len, kv_start_loc):
-        kv_indices[start : start + sl] = req_to_token_indexs[b][:sl]
-    kv_last_page_len_buffer = torch.empty(batch_size, device="cuda:0", dtype=torch.int32)
-    wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper(
-        workspace_buffer,
-        qo_indptr_buf=q_indptr,
-        paged_kv_indptr_buf=kv_indptr,
-        paged_kv_indices_buf=kv_indices,
-        paged_kv_last_page_len_buf=kv_last_page_len_buffer,
-    )
-    kv_last_page_len = torch.full((batch_size,), page_size, dtype=torch.int32)
-    k_cache = kv[:, :KV_HEADS, :].contiguous()
-    v_cache = kv[:, KV_HEADS:, :].contiguous()
-    k, k_scale = scaled_fp8_quant(k_cache.view(1, -1))
-    v, v_scale = scaled_fp8_quant(v_cache.view(1, -1))
-    wrapper.plan(
-        q_indptr,
-        kv_indptr,
-        kv_indices,
-        kv_last_page_len,
-        q_heads,
-        kv_heads,
-        head_dim,
-        page_size,
-        causal=True,
-        pos_encoding_mode="NONE",
-        logits_soft_cap=0.0,
-        q_data_type=q.dtype,
-        kv_data_type=torch.float8_e4m3fn,
-    )
-    wrapper.run(
-        q,
-        (k.view(-1, 1, kv_heads, head_dim), v.view(-1, 1, kv_heads, head_dim)),
-        k_scale=k_scale,
-        v_scale=v_scale,
-        out=o1,
-        return_lse=False,
-    )
-
-    # assert torch.allclose(o, o1, atol=1e-2, rtol=2e-1)
-    cos_sim1 = F.cosine_similarity(o, o1).mean()
-    print(cos_sim1)
-    assert cos_sim1 == 1
-
-
-if __name__ == "__main__":
-    test_context_attention_fwd_flashinfer_fp8(16, 1024, 28, 4, 128)
diff --git a/unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse_stage2.py b/unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse_stage2.py
deleted file mode 100644
index 2ba085cc91..0000000000
--- a/unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse_stage2.py
+++ /dev/null
@@ -1,129 +0,0 @@
-import pytest
-import torch
-from lightllm.utils.light_utils import light_ops
-
-
-def create_tensors(shared_seq_len):
-    batch_size = 4
-    num_heads = 32
-    kv_head_num = 8
-    seq_len = 256
-    head_dim = 128
-    max_len_in_batch = seq_len
-    block_seq = 256
-    max_batch_group_size = 4
-    quant_group_size = 8
-
-    test_dtype = torch.bfloat16
-
-    kv_shape = (batch_size * seq_len, kv_head_num, head_dim)
-    kv_scale_shape = (batch_size * seq_len, kv_head_num, head_dim // quant_group_size)
-
-    q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device="cuda")
-    k = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda")
-    k_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda")
-    v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda")
-    v_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda")
-    Req_to_tokens = torch.arange(0, seq_len * batch_size, dtype=torch.int32, device="cuda").view(batch_size, seq_len)
-    B_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda")
-    b_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda")
-    b_shared_seq_len = torch.full((batch_size,), shared_seq_len, dtype=torch.int32, device="cuda")
-    b_mark_shared_group = torch.ones(batch_size, dtype=torch.int32, device="cuda")
-    mid_out = torch.zeros(
-        size=(batch_size, num_heads, (seq_len // block_seq) + 2, head_dim), dtype=q.dtype, device="cuda"
-    )
-    mid_out_logsumexp = torch.zeros(
-        size=(batch_size, num_heads, (seq_len // block_seq) + 2), dtype=q.dtype, device="cuda"
-    )
-
-    return {
-        "q": q,
-        "k": k,
-        "k_scale": k_scale,
-        "v": v,
-        "v_scale": v_scale,
-        "Req_to_tokens": Req_to_tokens,
-        "B_req_idx": B_req_idx,
-        "b_seq_len": b_seq_len,
-        "b_shared_seq_len": b_shared_seq_len,
-        "b_mark_shared_group": b_mark_shared_group,
-        "max_len_in_batch": max_len_in_batch,
-        "mid_out": mid_out,
-        "mid_out_logsumexp": mid_out_logsumexp,
-        "block_seq": block_seq,
-        "max_batch_group_size": max_batch_group_size,
-        "head_dim": head_dim,
-    }
-
-
-@pytest.mark.parametrize("shared_seq_len", [0, 47, 77, 128, 200, 255])
-def test_flash_decode_stage2_execution(shared_seq_len):
-    setup_tensors = create_tensors(shared_seq_len)
-
-    light_ops.group8_int8kv_flashdecoding_diverse_stage2(
-        setup_tensors["block_seq"],
-        setup_tensors["mid_out"],
-        setup_tensors["mid_out_logsumexp"],
-        1.0 / (setup_tensors["head_dim"] ** 0.5),
-        setup_tensors["q"],
-        setup_tensors["k"],
-        setup_tensors["k_scale"],
-        setup_tensors["v"],
-        setup_tensors["v_scale"],
-        setup_tensors["Req_to_tokens"],
-        setup_tensors["B_req_idx"],
-        setup_tensors["b_seq_len"],
-        setup_tensors["b_shared_seq_len"],
-        setup_tensors["max_len_in_batch"],
-    )
-    seq_block_idx = (setup_tensors["b_shared_seq_len"][0].item() + setup_tensors["block_seq"] - 1) // setup_tensors[
-        "block_seq"
-    ]
-    mid_out = setup_tensors["mid_out"][:, :, seq_block_idx:, :]
-    mid_out_logsumexp = setup_tensors["mid_out_logsumexp"][:, :, seq_block_idx:]
-
-    q = setup_tensors["q"]
-    k = setup_tensors["k"]
-    v = setup_tensors["v"]
-    true_mid_out = torch.zeros_like(mid_out)
-    true_mid_out_logsumexp = torch.zeros_like(mid_out_logsumexp)
-    new_q = q
-    new_k = k.to(q.dtype)
-    new_v = v.to(q.dtype)
-
-    b_seq_len = setup_tensors["b_seq_len"] - setup_tensors["b_shared_seq_len"]
-    req_to_tokens = setup_tensors["Req_to_tokens"][:, setup_tensors["b_shared_seq_len"][0].item() :]
-
-    from lightllm.models.llama.triton_kernel.gqa_flash_decoding_stage1 import (
-        flash_decode_stage1 as gqa_flash_decode_stage1,
-    )
-
-    gqa_flash_decode_stage1(
-        q=new_q,
-        k=new_k,
-        v=new_v,
-        Req_to_tokens=req_to_tokens,
-        B_req_idx=setup_tensors["B_req_idx"],
-        B_Seqlen=b_seq_len,
-        max_len_in_batch=setup_tensors["max_len_in_batch"],
-        mid_out=true_mid_out,
-        mid_out_logsumexp=true_mid_out_logsumexp,
-        block_seq=setup_tensors["block_seq"],
-    )
-    print(f"\nshared_seq_len={shared_seq_len}")
-    print(f"mid_out: {mid_out[0:4, 0, 0, 0]}")
-    print(f"true_mid_out: {true_mid_out[0:4, 0, 0, 0]}")
-    abs_diff = (mid_out - true_mid_out).abs()
-    max_diff = abs_diff.max()
-    max_diff_idx = abs_diff.argmax()
-    max_diff_idx_unraveled = torch.unravel_index(max_diff_idx, abs_diff.shape)
-    mid_out_value = mid_out[max_diff_idx_unraveled]
-    true_mid_out_value = true_mid_out[max_diff_idx_unraveled]
-    print(f"max abs diff: {max_diff}, mid_out value: {mid_out_value}, " f"true_mid_out value: {true_mid_out_value}")
-
-    assert torch.allclose(
-        mid_out[0:4, 0, 0, 0], true_mid_out[0:4, 0, 0, 0], atol=1e-2
-    ), f"Mid output does not match expected values for shared_seq_len={shared_seq_len}"
-    assert torch.allclose(
-        mid_out_logsumexp, true_mid_out_logsumexp, atol=1e-2
-    ), f"LogSumExp output does not match expected values for shared_seq_len={shared_seq_len}"
diff --git a/unit_tests/models/llama/test_token_attention_nopad.py b/unit_tests/models/llama/test_token_attention_nopad.py
deleted file mode 100644
index 1bbb291662..0000000000
--- a/unit_tests/models/llama/test_token_attention_nopad.py
+++ /dev/null
@@ -1,151 +0,0 @@
-import torch
-import time
-import pytest
-import numpy as np
-import torch.nn.functional as F
-import flashinfer
-from lightllm.utils.log_utils import init_logger
-from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.common.req_manager import ReqManager
-from lightllm.models.llama.triton_kernel.gqa_decode_flashattention_nopad import gqa_decode_attention_fwd
-
-logger = init_logger(__name__)
-
-seed = 42
-torch.manual_seed(seed)
-
-if torch.cuda.is_available():
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-
-def ref_token_attention_nopad(q, k, v, o, q_h, h_dim, infer_state):
-    from lightllm.models.llama.triton_kernel.token_attention_nopad_att1 import token_att_fwd
-
-    total_token_num = infer_state.total_token_num
-    batch_size = infer_state.batch_size
-    calcu_shape1 = (batch_size, q_h, h_dim)
-
-    att_m_tensor = torch.empty((q_h, total_token_num), dtype=torch.float32).cuda()
-
-    token_att_fwd(
-        q.view(calcu_shape1),
-        k,
-        att_m_tensor,
-        infer_state.req_manager.req_to_token_indexs,
-        infer_state.b_req_idx,
-        infer_state.b_start_loc,
-        infer_state.b_seq_len,
-        infer_state.max_len_in_batch,
-    )
-
-    from lightllm.models.llama.triton_kernel.token_attention_softmax_and_reducev import (
-        token_softmax_reducev_fwd,
-    )
-
-    token_softmax_reducev_fwd(
-        att_m_tensor,
-        v,
-        o,
-        infer_state.req_manager.req_to_token_indexs,
-        infer_state.b_req_idx,
-        infer_state.b_start_loc,
-        infer_state.b_seq_len,
-    )
-    return o
-
-
-@pytest.mark.parametrize(
-    "batch, seqlen, q_heads, kv_heads, head_dim",
-    [
-        (a, b, c, d, e)
-        for a in [1, 16, 32, 128, 512]
-        for b in [16, 32, 512, 1024]
-        for c in [28]
-        for d in [4]
-        for e in [128]
-    ],
-)
-def test_token_attention_nopad(batch, seqlen, q_heads, kv_heads, head_dim):
-    Z, N_CTX, Q_HEADS, KV_HEADS, HEAD_DIM = batch, seqlen, q_heads, kv_heads, head_dim
-    dtype = torch.bfloat16
-    q = torch.randn((Z, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    kv = torch.randn((Z * N_CTX, 2 * KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-
-    max_input_len = Z * N_CTX
-    req_to_token_indexs = torch.randperm(max_input_len, dtype=torch.int32).cuda().view(Z, N_CTX)
-    b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * N_CTX
-    b_start_loc = torch.arange(Z).cuda().int() * N_CTX
-    b_req_idx = torch.randperm(Z, dtype=torch.int32).cuda()
-
-    o = torch.zeros((Z, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    o1 = torch.zeros((Z, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-
-    infer_state = LlamaInferStateInfo()
-    infer_state.batch_size = Z
-    infer_state.max_len_in_batch = N_CTX
-    infer_state.total_token_num = Z * N_CTX
-    infer_state.req_manager = ReqManager(Z, N_CTX, None)
-    infer_state.req_manager.req_to_token_indexs = req_to_token_indexs
-    infer_state.b_req_idx = b_req_idx
-    infer_state.b_seq_len = b_seq_len
-    infer_state.b_start_loc = b_start_loc
-
-    ref_token_attention_nopad(
-        q,
-        kv[:, :KV_HEADS, :],
-        kv[:, KV_HEADS:, :],
-        o,
-        Q_HEADS,
-        HEAD_DIM,
-        infer_state,
-    )
-    # gqa_decode_attention_fwd(
-    #     q,
-    #     kv[:,:KV_HEADS,:],
-    #     kv[:,KV_HEADS:,:],
-    #     o,
-    #     infer_state.req_manager.req_to_token_indexs,
-    #     infer_state.b_req_idx,
-    #     infer_state.b_seq_len,
-    # )
-
-    batch_size = Z
-    head_dim = HEAD_DIM
-    q_heads = Q_HEADS
-    kv_heads = KV_HEADS
-    page_size = 1
-    workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8).to(0)
-    kv_starts = torch.zeros((Z + 1,)).int().cuda()
-    kv_starts[1:] = torch.cumsum(b_seq_len, dim=0)
-    kv_indptr = kv_starts
-    kv_indices = torch.arange(Z * N_CTX).cuda().int()
-    for b, sl, start in zip(b_req_idx, b_seq_len, b_start_loc):
-        kv_indices[start : start + sl] = req_to_token_indexs[b][:sl]
-    kv_last_page_len_buffer = torch.empty(batch_size, device="cuda:0", dtype=torch.int32)
-    wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(
-        workspace_buffer,
-        "NHD",
-        use_cuda_graph=True,
-        use_tensor_cores=True,
-        paged_kv_indptr_buffer=kv_indptr,
-        paged_kv_indices_buffer=kv_indices,
-        paged_kv_last_page_len_buffer=kv_last_page_len_buffer,
-    )
-    kv_last_page_len_buffer = torch.full((batch_size,), page_size, dtype=torch.int32)
-    wrapper.plan(
-        kv_indptr,
-        kv_indices,
-        kv_last_page_len_buffer,
-        q_heads,
-        kv_heads,
-        head_dim,
-        page_size,
-        q_data_type=dtype,
-        non_blocking=True,
-    )
-    kv = kv.unsqueeze(1)
-    wrapper.run(q, (kv[:, :, :KV_HEADS, :], kv[:, :, KV_HEADS:, :]), out=o1, return_lse=False)
-
-    cos_sim1 = F.cosine_similarity(o, o1).mean()
-    assert cos_sim1 == 1.0
diff --git a/unit_tests/models/llama/test_token_attention_nopad_fa3_fp8.py b/unit_tests/models/llama/test_token_attention_nopad_fa3_fp8.py
deleted file mode 100644
index a7f48ab899..0000000000
--- a/unit_tests/models/llama/test_token_attention_nopad_fa3_fp8.py
+++ /dev/null
@@ -1,187 +0,0 @@
-import torch
-import time
-import pytest
-import numpy as np
-import torch.nn.functional as F
-from lightllm.utils.log_utils import init_logger
-from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.models.llama.triton_kernel.gqa_decode_flashattention_nopad import gqa_decode_attention_fwd
-from lightllm.utils.sgl_utils import flash_attn_with_kvcache
-from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
-
-logger = init_logger(__name__)
-
-seed = 42
-torch.manual_seed(seed)
-
-if torch.cuda.is_available():
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-
-def kv_quantize_per_head_fp8(kv_buffer: torch.Tensor, seq_lens):
-    device = kv_buffer.device
-    B = seq_lens.size(0)
-    min_fp8 = torch.finfo(torch.float8_e4m3fn).min
-    max_fp8 = torch.finfo(torch.float8_e4m3fn).max
-    _, S_max, H, D = kv_buffer.shape
-    seq_range = torch.arange(S_max, device=device)[None, :]
-    valid_mask = (seq_range < seq_lens[:, None]).view(B, S_max, 1, 1)
-    masked = kv_buffer * valid_mask
-    max_per_bh = masked.float().abs().amax(dim=(1, 3))  # [B, H]
-    scales = torch.where(max_per_bh > 0, max_per_bh / max_fp8, torch.ones_like(max_per_bh))
-    scales_exp = scales.view(B, 1, H, 1)
-    q = (kv_buffer / scales_exp).clamp(min_fp8, max_fp8).to(torch.float8_e4m3fn)
-    return q, scales
-
-
-def ref_token_attention_nopad(q, k, v, o, q_h, h_dim, infer_state, req_to_token_indexs):
-    from lightllm.models.llama.triton_kernel.token_attention_nopad_att1 import token_att_fwd
-
-    total_token_num = infer_state.total_token_num
-    batch_size = infer_state.batch_size
-    calcu_shape1 = (batch_size, q_h, h_dim)
-
-    att_m_tensor = torch.empty((q_h, total_token_num), dtype=torch.float32).cuda()
-
-    token_att_fwd(
-        q.view(calcu_shape1),
-        k,
-        att_m_tensor,
-        req_to_token_indexs,
-        infer_state.b_req_idx,
-        infer_state.b_start_loc,
-        infer_state.b_seq_len,
-        infer_state.max_len_in_batch,
-    )
-
-    from lightllm.models.llama.triton_kernel.token_attention_softmax_and_reducev import (
-        token_softmax_reducev_fwd,
-    )
-
-    token_softmax_reducev_fwd(
-        att_m_tensor,
-        v,
-        o,
-        req_to_token_indexs,
-        infer_state.b_req_idx,
-        infer_state.b_start_loc,
-        infer_state.b_seq_len,
-    )
-    return o
-
-
-@pytest.mark.parametrize(
-    "batch, seqlen, q_heads, kv_heads, head_dim",
-    [
-        (a, b, c, d, e)
-        for a in [1, 16, 32, 128, 512]
-        for b in [16, 32, 512, 1024]
-        for c in [28]
-        for d in [4]
-        for e in [128]
-    ],
-)
-def test_token_attention_nopad_fa3_fp8(batch, seqlen, q_heads, kv_heads, head_dim):
-    Z, N_CTX, Q_HEADS, KV_HEADS, HEAD_DIM = batch, seqlen, q_heads, kv_heads, head_dim
-    dtype = torch.bfloat16
-    q = torch.randn((Z, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    kv = torch.randn((Z * N_CTX, 2 * KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    # for i in range(Z * N_CTX):
-    #     kv[i] = torch.randn((2 * KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda") * (i % 10 + 1)
-
-    max_input_len = Z * N_CTX
-    req_to_token_indexs = torch.randperm(max_input_len, dtype=torch.int32).cuda().view(Z, N_CTX)
-    b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * (N_CTX // 2)
-    rand_num = torch.randint_like(b_seq_len, high=(N_CTX // 2), dtype=torch.int32, device="cuda")
-    b_seq_len += rand_num
-    b_start_loc = b_seq_len.cumsum(0) - b_seq_len
-    b_req_idx = torch.randperm(Z, dtype=torch.int32).cuda()
-
-    o = torch.zeros((Z, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    o1 = torch.zeros((Z, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-
-    infer_state = LlamaInferStateInfo()
-    infer_state.batch_size = Z
-    infer_state.max_len_in_batch = N_CTX
-    infer_state.total_token_num = Z * N_CTX
-    infer_state.b_req_idx = b_req_idx
-    infer_state.b_seq_len = b_seq_len
-    infer_state.b_start_loc = b_start_loc
-
-    ref_token_attention_nopad(
-        q,
-        kv[:, :KV_HEADS, :],
-        kv[:, KV_HEADS:, :],
-        o,
-        Q_HEADS,
-        HEAD_DIM,
-        infer_state,
-        req_to_token_indexs,
-    )
-    # gqa_decode_attention_fwd(
-    #     q,
-    #     kv[:,:KV_HEADS,:],
-    #     kv[:,KV_HEADS:,:],
-    #     o,
-    #     req_to_token_indexs,
-    #     infer_state.b_req_idx,
-    #     infer_state.b_seq_len,
-    # )
-
-    batch_size = Z
-    head_dim = HEAD_DIM
-    q_heads = Q_HEADS
-    kv_heads = KV_HEADS
-    kv_starts = torch.zeros((Z + 1,)).int().cuda()
-    kv_starts[1:] = torch.cumsum(b_seq_len, dim=0)
-    q_starts = torch.arange(0, Z + 1).int().cuda()
-    page_table = torch.empty((batch_size, N_CTX), dtype=torch.int32).to(0)
-    page_table.copy_(req_to_token_indexs[b_req_idx, :N_CTX])
-
-    k_cache = kv[:, :KV_HEADS, :].contiguous()
-    v_cache = kv[:, KV_HEADS:, :].contiguous()
-    # o1 = flash_attn_with_kvcache(
-    #     q=q,
-    #     k_cache=k_cache[page_table].view(-1, N_CTX, kv_heads, head_dim),
-    #     v_cache=v_cache[page_table].view(-1, N_CTX, kv_heads, head_dim),
-    #     # page_table=page_table,
-    #     cache_seqlens=infer_state.b_seq_len,
-    #     cu_seqlens_q=q_starts,
-    #     cu_seqlens_k_new=kv_starts,
-    #     max_seqlen_q=1,
-    #     causal=False,
-    #     window_size=(-1, -1),
-    #     softcap=0.0,
-    #     return_softmax_lse=False,
-    # )
-
-    q, q_scale = scaled_fp8_quant(q.view(batch_size * kv_heads, -1), use_per_token_if_dynamic=True)
-    k, k_scale = kv_quantize_per_head_fp8(k_cache[page_table], b_seq_len)
-    v, v_scale = kv_quantize_per_head_fp8(v_cache[page_table], b_seq_len)
-    o1 = flash_attn_with_kvcache(
-        q=q.view(-1, q_heads, head_dim),
-        k_cache=k.view(-1, N_CTX, kv_heads, head_dim),
-        v_cache=v.view(-1, N_CTX, kv_heads, head_dim),
-        # page_table=page_table,
-        cache_seqlens=infer_state.b_seq_len,
-        cu_seqlens_q=q_starts,
-        cu_seqlens_k_new=kv_starts,
-        max_seqlen_q=1,
-        causal=False,
-        window_size=(-1, -1),
-        softcap=0.0,
-        q_descale=q_scale.view(batch_size, kv_heads),
-        k_descale=k_scale.view(batch_size, kv_heads),
-        v_descale=v_scale.view(batch_size, kv_heads),
-        return_softmax_lse=False,
-    )
-
-    # assert torch.allclose(o, o1, atol=1e-1, rtol=1e-1)
-    cos_sim1 = F.cosine_similarity(o, o1).mean()
-    print(cos_sim1)
-    assert cos_sim1 == 1
-
-
-if __name__ == "__main__":
-    test_token_attention_nopad_fa3_fp8(16, 16384, 28, 4, 128)
diff --git a/unit_tests/models/llama/test_token_attention_nopad_flashinfer_fp8.py b/unit_tests/models/llama/test_token_attention_nopad_flashinfer_fp8.py
deleted file mode 100644
index 5c0e595b96..0000000000
--- a/unit_tests/models/llama/test_token_attention_nopad_flashinfer_fp8.py
+++ /dev/null
@@ -1,170 +0,0 @@
-import torch
-import time
-import pytest
-import numpy as np
-import torch.nn.functional as F
-import flashinfer
-from lightllm.utils.log_utils import init_logger
-from lightllm.models.llama.infer_struct import LlamaInferStateInfo
-from lightllm.models.llama.triton_kernel.gqa_decode_flashattention_nopad import gqa_decode_attention_fwd
-from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
-
-logger = init_logger(__name__)
-
-seed = 42
-torch.manual_seed(seed)
-
-if torch.cuda.is_available():
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-
-def ref_token_attention_nopad(q, k, v, o, q_h, h_dim, infer_state, req_to_token_indexs):
-    from lightllm.models.llama.triton_kernel.token_attention_nopad_att1 import token_att_fwd
-
-    total_token_num = infer_state.total_token_num
-    batch_size = infer_state.batch_size
-    calcu_shape1 = (batch_size, q_h, h_dim)
-
-    att_m_tensor = torch.empty((q_h, total_token_num), dtype=torch.float32).cuda()
-
-    token_att_fwd(
-        q.view(calcu_shape1),
-        k,
-        att_m_tensor,
-        req_to_token_indexs,
-        infer_state.b_req_idx,
-        infer_state.b_start_loc,
-        infer_state.b_seq_len,
-        infer_state.max_len_in_batch,
-    )
-
-    from lightllm.models.llama.triton_kernel.token_attention_softmax_and_reducev import (
-        token_softmax_reducev_fwd,
-    )
-
-    token_softmax_reducev_fwd(
-        att_m_tensor,
-        v,
-        o,
-        req_to_token_indexs,
-        infer_state.b_req_idx,
-        infer_state.b_start_loc,
-        infer_state.b_seq_len,
-    )
-    return o
-
-
-@pytest.mark.parametrize(
-    "batch, seqlen, q_heads, kv_heads, head_dim",
-    [
-        (a, b, c, d, e)
-        for a in [1, 16, 32, 128, 512]
-        for b in [16, 32, 512, 1024]
-        for c in [28]
-        for d in [4]
-        for e in [128]
-    ],
-)
-def test_token_attention_nopad_flashinfer_fp8(batch, seqlen, q_heads, kv_heads, head_dim):
-    Z, N_CTX, Q_HEADS, KV_HEADS, HEAD_DIM = batch, seqlen, q_heads, kv_heads, head_dim
-    dtype = torch.bfloat16
-    q = torch.randn((Z, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    kv = torch.randn((Z * N_CTX, 2 * KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    # for i in range(Z * N_CTX):
-    #     kv[i] = torch.randn((2 * KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda") * (i % 10 + 1)
-
-    max_input_len = Z * N_CTX
-    req_to_token_indexs = torch.randperm(max_input_len, dtype=torch.int32).cuda().view(Z, N_CTX)
-    b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * (N_CTX // 2)
-    rand_num = torch.randint_like(b_seq_len, high=(N_CTX // 2), dtype=torch.int32, device="cuda")
-    b_seq_len += rand_num
-    b_start_loc = b_seq_len.cumsum(0) - b_seq_len
-    b_req_idx = torch.randperm(Z, dtype=torch.int32).cuda()
-
-    o = torch.zeros((Z, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-    o1 = torch.zeros((Z, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda")
-
-    infer_state = LlamaInferStateInfo()
-    infer_state.batch_size = Z
-    infer_state.max_len_in_batch = N_CTX
-    infer_state.total_token_num = Z * N_CTX
-    infer_state.b_req_idx = b_req_idx
-    infer_state.b_seq_len = b_seq_len
-    infer_state.b_start_loc = b_start_loc
-
-    ref_token_attention_nopad(
-        q,
-        kv[:, :KV_HEADS, :],
-        kv[:, KV_HEADS:, :],
-        o,
-        Q_HEADS,
-        HEAD_DIM,
-        infer_state,
-        req_to_token_indexs,
-    )
-    # gqa_decode_attention_fwd(
-    #     q,
-    #     kv[:,:KV_HEADS,:],
-    #     kv[:,KV_HEADS:,:],
-    #     o,
-    #     req_to_token_indexs,
-    #     infer_state.b_req_idx,
-    #     infer_state.b_seq_len,
-    # )
-
-    batch_size = Z
-    head_dim = HEAD_DIM
-    q_heads = Q_HEADS
-    kv_heads = KV_HEADS
-    page_size = 1
-    workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8).to(0)
-    kv_starts = torch.zeros((Z + 1,)).int().cuda()
-    kv_starts[1:] = torch.cumsum(b_seq_len, dim=0)
-    kv_indptr = kv_starts
-    kv_indices = torch.arange(Z * N_CTX).cuda().int()
-    for b, sl, start in zip(b_req_idx, b_seq_len, b_start_loc):
-        kv_indices[start : start + sl] = req_to_token_indexs[b][:sl]
-    kv_last_page_len_buffer = torch.empty(batch_size, device="cuda:0", dtype=torch.int32)
-    wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(
-        workspace_buffer,
-        "NHD",
-        use_cuda_graph=True,
-        use_tensor_cores=True,
-        paged_kv_indptr_buffer=kv_indptr,
-        paged_kv_indices_buffer=kv_indices,
-        paged_kv_last_page_len_buffer=kv_last_page_len_buffer,
-    )
-    kv_last_page_len_buffer = torch.full((batch_size,), page_size, dtype=torch.int32)
-    k_cache = kv[:, :KV_HEADS, :].contiguous()
-    v_cache = kv[:, KV_HEADS:, :].contiguous()
-    k, k_scale = scaled_fp8_quant(k_cache.view(1, -1))
-    v, v_scale = scaled_fp8_quant(v_cache.view(1, -1))
-    wrapper.plan(
-        kv_indptr,
-        kv_indices,
-        kv_last_page_len_buffer,
-        q_heads,
-        kv_heads,
-        head_dim,
-        page_size,
-        q_data_type=dtype,
-        kv_data_type=torch.float8_e4m3fn,
-        non_blocking=True,
-    )
-    wrapper.run(
-        q,
-        (k.view(-1, 1, kv_heads, head_dim), v.view(-1, 1, kv_heads, head_dim)),
-        k_scale=k_scale,
-        v_scale=v_scale,
-        out=o1,
-        return_lse=False,
-    )
-
-    cos_sim1 = F.cosine_similarity(o, o1).mean()
-    print(cos_sim1)
-    assert cos_sim1 == 1.0
-
-
-if __name__ == "__main__":
-    test_token_attention_nopad_flashinfer_fp8(16, 16384, 28, 4, 128)
diff --git a/unit_tests/models/qwen3-vl/test_deepstack_emb.py b/unit_tests/models/qwen3-vl/test_deepstack_emb.py
index 2f929fe0d0..f629a16352 100644
--- a/unit_tests/models/qwen3-vl/test_deepstack_emb.py
+++ b/unit_tests/models/qwen3-vl/test_deepstack_emb.py
@@ -50,7 +50,7 @@ def test_deepstack_same_image_twice():
         deepstack_embs=deepstack_embs,
         img_token_lens=img_token_lens,
         img_start_token_ids=img_start_token_ids,
-        img_start_locs=img_start_locs,
+        img_start_locs_in_cache=img_start_locs,
     )
 
     # 7. 看看相同图片两段上的增量
diff --git a/unit_tests/server/core/objs/test_req.py b/unit_tests/server/core/objs/test_req.py
index 45fa7967fb..1c946531c1 100644
--- a/unit_tests/server/core/objs/test_req.py
+++ b/unit_tests/server/core/objs/test_req.py
@@ -1,6 +1,22 @@
 import pytest
-
+import easydict
 from lightllm.server.core.objs.req import Req, TokenHealingReq, ChunkedPrefillReq, SamplingParams
+from lightllm.utils.envs_utils import set_env_start_args
+
+
+@pytest.fixture(scope="module", autouse=True)
+def setup_module_env():
+    set_env_start_args(
+        easydict.EasyDict(
+            {
+                "mtp_step": 0,
+                "llm_prefill_att_backend": ["None"],
+                "llm_decode_att_backend": ["None"],
+                "cpu_cache_token_page_size": 256,
+                "enable_cpu_cache": False,
+            }
+        )
+    )
 
 
 @pytest.fixture
diff --git a/unit_tests/server/core/objs/test_shm_req_manager.py b/unit_tests/server/core/objs/test_shm_req_manager.py
index 1d1ae2ef1a..e26f128d5b 100644
--- a/unit_tests/server/core/objs/test_shm_req_manager.py
+++ b/unit_tests/server/core/objs/test_shm_req_manager.py
@@ -14,8 +14,11 @@ def setup_env():
             running_max_req_size=10,
             disable_chunked_prefill=True,
             token_healing_mode=False,
-            enable_flashinfer_prefill=False,
-            enable_flashinfer_decode=False,
+            mtp_step=0,
+            llm_prefill_att_backend=["None"],
+            llm_decode_att_backend=["None"],
+            cpu_cache_token_page_size=256,
+            enable_cpu_cache=False,
         )
     )
     # clear the lru_cache if used