[Qwen-moe] use npu_add_rms_norm_quant operator

Signed-off-by: s30076806 <[email protected]>

Author: s30076806

vllm_ascend/envs.py

@@ -159,6 +159,11 @@
     #   1: enable moe all2all seq.
     "VLLM_ASCEND_ENABLE_MOE_ALL2ALL_SEQ":
     lambda: bool(int(os.getenv('VLLM_ASCEND_ENABLE_MOE_ALL2ALL_SEQ', '0'))),
+    # Whether to enable layernorm.py torch_npu.npu_add_rms_norm_quant
+    # 0: default
+    # 1: enable
+    "USE_ADD_RMSNORM_QUANT":
+    lambda: int(os.getenv("USE_ADD_RMSNORM_QUANT", '0')),
 }
 
 # end-env-vars-definition

vllm_ascend/models/qwen3_moe.py

@@ -50,6 +50,10 @@
 from vllm_ascend.ops.fused_moe import AscendFusedMoE
 from vllm_ascend.ops.sequence_parallel import (MetadataForPadding,
                                                init_metadata_for_sp)
+from vllm_ascend.ops.layernorm import AddRMSNormW8A8Quant
+from vllm_ascend.quantization.w8a8 import AscendW8A8LinearMethod
+from vllm_ascend.quantization.quant_config import AscendQuantConfig
+import vllm_ascend.envs as envs
 
 
 class CustomSparseMoeBlock(Qwen3MoeSparseMoeBlock):
@@ -183,8 +187,21 @@ def __init__(
                                    hidden_act=config.hidden_act,
                                    quant_config=quant_config,
                                    prefix=f"{prefix}.mlp")
-        self.input_layernorm = RMSNorm(config.hidden_size,
-                                       eps=config.rms_norm_eps)
+        if not envs.USE_ADD_RMSNORM_QUANT:
+            self.input_layernorm = RMSNorm(config.hidden_size,
+                                           eps=config.rms_norm_eps)
+        else:
+            assert isinstance(quant_config, AscendQuantConfig), \
+                "Expected quant_config to be an instance of AscendQuantConfig"
+            if isinstance(self.self_attn.qkv_proj.quant_method.quant_method,
+                          AscendW8A8LinearMethod):
+                self.input_layernorm = AddRMSNormW8A8Quant(
+                    config.hidden_size,
+                    layer=self.self_attn.qkv_proj,
+                    eps=config.rms_norm_eps)
+            else:
+                self.input_layernorm = RMSNorm(config.hidden_size,
+                                               eps=config.rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(config.hidden_size,
                                                 eps=config.rms_norm_eps)