[main] addrmsnorm + quant fusion optim

Signed-off-by: rjg-lyh <1318825571@qq.com>

Author: rjg-lyh

vllm_ascend/ascend_forward_context.py

@@ -66,7 +66,8 @@ def set_ascend_forward_context(
         moe_comm_method: str = "",
         num_actual_tokens: Optional[int] = None,
         aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
-        batch_descriptor: Optional[BatchDescriptor] = None):
+        batch_descriptor: Optional[BatchDescriptor] = None,
+        prefetch_model: torch.nn.Module = None):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.
     We add some additional param into forward_context.
@@ -81,6 +82,19 @@ def set_ascend_forward_context(
             batch_descriptor=batch_descriptor,
     ):
         forward_context = get_forward_context()
+        if envs_ascend.VLLM_ASCEND_ENABLE_ADDRMSNORM_QUANT_FUSION:
+            model_type = vllm_config.model_config.hf_config.model_type
+            forward_context.prefetch_model = prefetch_model
+            forward_context.layer_idx = 0
+            # dense model
+            if model_type in ["llama", "qwen2", "qwen3"]:
+                forward_context.fusion_linear = "gate_up_dense"
+            # moe model
+            elif model_type == "qwen3_moe":
+                forward_context.fusion_linear = "qkv_moe"
+            else:
+                raise ValueError(f"AddRmsNorm+Quant Fusion unsupport model type: {model_type}")
+                
         forward_context.moe_comm_method_name = moe_comm_method + "commimpl"
         forward_context.with_prefill = with_prefill
         tp_world_size = get_tensor_model_parallel_world_size()

vllm_ascend/envs.py

@@ -144,6 +144,9 @@
     # this feature in eager mode will get better performance.
     "VLLM_ASCEND_ENABLE_MLP_OPTIMIZE":
     lambda: bool(int(os.getenv("VLLM_ASCEND_ENABLE_MLP_OPTIMIZE", '0'))),
+    # Whether to enable addrmsnorm + quant fusion
+    "VLLM_ASCEND_ENABLE_ADDRMSNORM_QUANT_FUSION":
+    lambda: bool(int(os.getenv("VLLM_ASCEND_ENABLE_ADDRMSNORM_QUANT_FUSION", '0'))),
     # Determine the number of physical devices in a non-full-use scenario
     # caused by the initialization of the Mooncake connector.
     "PHYSICAL_DEVICES":

vllm_ascend/ops/layernorm.py

@@ -18,7 +18,9 @@
 from typing import Optional, Tuple, Union
 
 import torch
+from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.layernorm import RMSNorm
+import vllm_ascend.envs as envs_ascend
 
 
 class AddRMSNormW8A8Quant(RMSNorm):
@@ -64,6 +66,42 @@ def forward(
                                              self.variance_epsilon)
         return x
 
+def _addrmsnorm_w8a8_quant_forward_oot(
+    self,
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    layer: torch.nn.Module,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    import torch_npu
+
+    x, _, residual = torch_npu.npu_add_rms_norm_quant(
+        x,
+        residual,
+        self.weight,
+        layer.aclnn_input_scale,
+        layer.aclnn_input_offset,
+        epsilon=self.variance_epsilon)
+    return x, residual
+
+
+def _addrmsnorm_forward_oot(
+    self,
+    x: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    import torch_npu
+
+    from vllm_ascend.utils import is_310p
+    if is_310p():
+        orig_dtype = residual.dtype
+        x = x + residual.to(x.dtype)
+        residual = x.to(orig_dtype)
+        x, _ = torch_npu.npu_rms_norm(x, self.weight,
+                                        self.variance_epsilon)
+    else:
+        x, _, residual = torch_npu.npu_add_rms_norm(
+            x, residual, self.weight, self.variance_epsilon)
+    return x, residual
 
 class AscendRMSNorm(RMSNorm):
 
@@ -90,8 +128,7 @@ def forward_oot(
                 x, _ = torch_npu.npu_rms_norm(x, self.weight,
                                               self.variance_epsilon)
             else:
-                x, _, residual = torch_npu.npu_add_rms_norm(
-                    x, residual, self.weight, self.variance_epsilon)
+                x, residual = _addrmsnorm_forward_oot(x, residual)
             return x, residual
 
         x, residual = torch_npu.npu_rms_norm(x, self.weight,

vllm_ascend/worker/model_runner_v1.py

@@ -1563,7 +1563,8 @@ def execute_model(
                     aclgraph_runtime_mode=aclgraph_runtime_mode,
                     batch_descriptor=batch_descriptor,
                     num_actual_tokens=scheduler_output.
-                    total_num_scheduled_tokens):
+                    total_num_scheduled_tokens,
+                    prefetch_model=self.model):
                 self.maybe_setup_kv_connector(scheduler_output)
 
                 hidden_states = self._generate_process_reqs_hidden_states(
@@ -2009,7 +2010,8 @@ def dummy_compute_logits(hidden_states):
                     moe_comm_method=moe_comm_method,
                     num_actual_tokens=0,
                     aclgraph_runtime_mode=aclgraph_runtime_mode,
-                    batch_descriptor=batch_descriptor):
+                    batch_descriptor=batch_descriptor,
+                    prefetch_model=self.model):
                 hidden_states = self._generate_dummy_run_hidden_states(
                     with_prefill, is_torchair_compile, input_ids, positions,
                     attn_metadata, num_tokens, intermediate_tensors,