[Performance][310P] Replace l2norm with RMSNorm operator on 310P (vllm-project#10408)

### What this PR does / why we need it?
This pull request optimizes L2 normalization for the 310P NPU by
introducing a custom l2norm_310p function that leverages the
npu_rms_norm kernel, replacing standard PyTorch normalization in several
FLA operators. However, the current implementation has a mathematical
discrepancy where the epsilon parameter is not scaled by the dimension
size, leading to potential numerical inaccuracies. Additionally, it
lacks a CPU fallback, which will cause crashes when run on non-NPU
devices.
### Does this PR introduce _any_ user-facing change?
None
### How was this patch tested?


- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@967c5c3

Signed-off-by: YangShuai52 <shuaiyang047@163.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: YangShuai52

tests/ut/_310p/ops/test_chunk_gated_delta_rule_310.py

@@ -1,8 +1,29 @@
+from unittest.mock import patch
+
+import pytest
 import torch
+import torch_npu
 
 from vllm_ascend._310p.ops.fla.chunk_gated_delta_rule import chunk_gated_delta_rule_pytorch
 
 
+def _cpu_rms_norm(x, weight, eps):
+    """CPU fallback for torch_npu.npu_rms_norm used by l2norm_310p on CPU runners."""
+    orig_dtype = x.dtype
+    x32 = x.float()
+    var = x32.pow(2).mean(-1, keepdim=True)
+    x32 = x32 * torch.rsqrt(var + eps)
+    out = (x32 * weight.float()).to(orig_dtype)
+    return out, None
+
+
+@pytest.fixture(autouse=True)
+def _mock_npu_rms_norm():
+    # conftest stubs npu_rms_norm with a bare MagicMock(); override with a CPU impl.
+    with patch.object(torch_npu, "npu_rms_norm", side_effect=_cpu_rms_norm, create=True):
+        yield
+
+
 def test_chunk_gated_delta_rule_310_output_shape_and_dtype():
     torch.manual_seed(0)
 

vllm_ascend/_310p/ops/fla/__init__.py

@@ -1,9 +1,11 @@
 from .chunk_gated_delta_rule import chunk_gated_delta_rule_310
 from .fused_gdn_gating import fused_gdn_gating_pytorch
 from .fused_recurrent_gated_delta_rule import fused_recurrent_gated_delta_rule_pytorch
+from .l2norm import l2norm_310p
 
 __all__ = [
     "fused_gdn_gating_pytorch",
     "fused_recurrent_gated_delta_rule_pytorch",
     "chunk_gated_delta_rule_310",
+    "l2norm_310p",
 ]

vllm_ascend/_310p/ops/fla/chunk_gated_delta_rule.py

@@ -20,6 +20,8 @@
 import torch
 import torch.nn.functional as F
 
+from vllm_ascend._310p.ops.fla.l2norm import l2norm_310p
+
 CHUNK_SIZE = 64
 
 
@@ -105,8 +107,8 @@ def _torch_chunk_gated_delta_rule_chunked(
     """
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
-        query = F.normalize(query, p=2, dim=-1, eps=1e-6).to(query.dtype)
-        key = F.normalize(key, p=2, dim=-1, eps=1e-6).to(key.dtype)
+        query = l2norm_310p(query)
+        key = l2norm_310p(key)
 
     query, key, value, beta, g = [
         x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
@@ -200,7 +202,7 @@ def _require_ascend_chunk_ops(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor)
 def _maybe_l2norm(x: torch.Tensor, enabled: bool) -> torch.Tensor:
     if not enabled:
         return x
-    return F.normalize(x.to(torch.float32), p=2, dim=-1, eps=1e-6).to(x.dtype)
+    return l2norm_310p(x)
 
 
 def _pad_bthd_to_chunk(

vllm_ascend/_310p/ops/fla/fused_recurrent_gated_delta_rule.py

@@ -1,11 +1,12 @@
 import torch
-import torch.nn.functional as F
+
+from vllm_ascend._310p.ops.fla.l2norm import l2norm_310p
 
 
 def _maybe_l2norm(x: torch.Tensor, enabled: bool) -> torch.Tensor:
     if not enabled:
         return x
-    return F.normalize(x, p=2, dim=-1, eps=1e-6).to(x.dtype)
+    return l2norm_310p(x)
 
 
 def _expand_to_hv(x: torch.Tensor, hv: int) -> torch.Tensor:

vllm_ascend/_310p/ops/fla/gdn_310.py

@@ -18,7 +18,6 @@
 
 
 import torch
-import torch.nn.functional as F
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
 from vllm.v1.attention.backend import AttentionMetadata  # type: ignore
@@ -27,6 +26,7 @@
 
 from vllm_ascend._310p.ops.fla.chunk_gated_delta_rule import chunk_gated_delta_rule_310
 from vllm_ascend._310p.ops.fla.fused_gdn_gating import fused_gdn_gating_pytorch
+from vllm_ascend._310p.ops.fla.l2norm import l2norm_310p
 from vllm_ascend.ascend_forward_context import _EXTRA_CTX
 from vllm_ascend.attention.utils import maybe_save_kv_layer_to_connector
 from vllm_ascend.compilation.acl_graph import get_draft_graph_params, get_graph_params
@@ -111,10 +111,6 @@ def _register_310_conv1d_buffer_replay(
     graph_params.conv1d_events[num_actual_tokens].append(None)
 
 
-def _l2norm(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
-    return F.normalize(x.to(torch.float32), p=2, dim=-1, eps=eps).to(x.dtype)
-
-
 def _flatten_state_indices(
     ssm_state_indices: torch.Tensor,
     cu_seqlens: torch.Tensor,
@@ -161,8 +157,8 @@ def npu_recurrent_gated_delta_rule_310(
     use_qk_l2norm_in_kernel: bool = True,
 ) -> torch.Tensor:
     if use_qk_l2norm_in_kernel:
-        q = _l2norm(q)
-        k = _l2norm(k)
+        q = l2norm_310p(q)
+        k = l2norm_310p(k)
 
     total_tokens = v.shape[1]
     flat_state_indices = _flatten_state_indices(ssm_state_indices, cu_seqlens, total_tokens)

vllm_ascend/_310p/ops/fla/l2norm.py

@@ -0,0 +1,43 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm-ascend project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# mypy: ignore-errors
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch_npu
+from vllm.model_executor.layers.fla.ops.utils import tensor_cache
+
+
+@tensor_cache
+def _l2norm_unit_weight(dim: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
+    # RMSNorm with weight 1/sqrt(dim) matches L2 norm: x / sqrt(sum(x^2)).
+    return torch.full((dim,), 1.0 / math.sqrt(dim), dtype=dtype, device=device)
+
+
+def l2norm_310p(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    """L2-normalize the last dimension using the 310P NPU RMSNorm kernel."""
+    orig_shape = x.shape
+    dim = x.shape[-1]
+    x_2d = x.reshape(-1, dim).contiguous()
+    weight = _l2norm_unit_weight(dim, x.dtype, x.device)
+    # RMSNorm: y = x / sqrt(mean(x^2) + eps_rms) * weight
+    # With weight=1/sqrt(dim), this equals x / sqrt(sum(x^2) + dim * eps_rms).
+    # L2 norm needs y = x / sqrt(sum(x^2) + eps), so eps_rms = eps / dim.
+    y, _ = torch_npu.npu_rms_norm(x_2d, weight, eps / dim)
+    return y.reshape(orig_shape)