update l2_norm op

Author: maojunx99

xllm/core/kernels/ops_api.cpp

@@ -756,6 +756,14 @@ void fused_indexer_k(FusedIndexerKParams& params) {
 #endif
 }
 
+torch::Tensor l2_norm(torch::Tensor& x, double eps) {
+#if defined(USE_NPU)
+  return npu::npu_l2norm_last_dim(x, eps);
+#else
+  NOT_IMPLEMENTED();
+#endif
+}
+
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
 moe_init_routing_v2(MoeInitRoutingV2Params& params) {
 #if defined(USE_NPU)

xllm/core/kernels/ops_api.h

@@ -94,6 +94,9 @@ void fused_indexer_q(FusedIndexerQParams& params);
 
 void fused_indexer_k(FusedIndexerKParams& params);
 
+// L2 normalization along the last dimension
+torch::Tensor l2_norm(torch::Tensor& x, double eps = 1e-6);
+
 // TODO: NPU moe_init_routing_v2 is equivalent to moe_gen_idx + moe_expand_input
 // (and token_count/cusum outputs) on other backends.
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>

xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp

@@ -450,8 +450,8 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
         {input_params.block_tables.select(1, 0)},
         last_recurrent_state.transpose(-1, -2).to(ssm_cache.dtype()));
   } else {
-    processed_q = l2norm(processed_q, -1, 1e-6);
-    processed_k = l2norm(processed_k, -1, 1e-6);
+    processed_q = xllm::kernel::l2_norm(processed_q, 1e-6);
+    processed_k = xllm::kernel::l2_norm(processed_k, 1e-6);
     torch::Tensor ssm_state_indices =
         attn_metadata.block_table.select(1, 0).contiguous();
     auto zero = torch::zeros({1},attn_metadata.q_seq_lens.options());