fix AddRMSNormW8A8Quant init error 'TypeError:wrapper_rmsnorm_init..init() takes 2 positional arguments but 6 were given'

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

Author: socrahow

vllm_ascend/models/gemma3.py

@@ -79,7 +79,7 @@ def __init__(
                       AscendW8A8LinearMethod):
             self.pre_feedforward_layernorm = AddRMSNormW8A8Quant(
                 config.hidden_size,
-                layer=self.self_attn.qkv_proj,
+                layer=self.mlp.gate_up_proj,
                 eps=config.rms_norm_eps
             )
 

vllm_ascend/ops/layernorm.py

@@ -33,7 +33,11 @@ def __init__(
         has_weight: bool = True,
         dtype: Optional[torch.dtype] = None,
     ) -> None:
-        super().__init__(hidden_size=hidden_size, eps=eps, var_hidden_size=var_hidden_size, has_weight=has_weight, dtype=dtype)
+        super().__init__(hidden_size=hidden_size,
+                         eps=eps,
+                         var_hidden_size=var_hidden_size,
+                         has_weight=has_weight,
+                         dtype=dtype)
         self.layer = layer
 
     def forward(
@@ -59,14 +63,12 @@ def forward(
 
 
 class AscendRMSNorm(RMSNorm):
-
     def 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 residual is not None:
             if is_310p():