bugfix: add compiler/version.info to CANN version detection path for CANN 9.0.

test/python_test/test_layer_norm_fwd.py

@@ -138,4 +138,28 @@ def test_layer_norm_fwd_npu(dtype, batch, full_n, group_size, has_bias,
     assert torch.allclose(rstd_npu, rstd_ref, atol=1e-2, rtol=1e-2)
     if not is_rms_norm:
         mean_npu = mean_npu.cpu().float()
-        assert torch.allclose(mean_npu, mean_ref, atol=1e-2, rtol=1e-2)
+        assert torch.allclose(mean_npu, mean_ref, atol=1e-2, rtol=1e-2)
+
+
+def test_layer_norm_fwd_large_n128_ub_budget():
+    try:
+        torch_npu.npu.set_device(0)
+    except Exception as e:
+        pytest.skip(f"NPU device not available: {e}")
+
+    torch.manual_seed(2026)
+    eps = 1e-6
+    x = torch.randn(4096, 128, dtype=torch.bfloat16)
+    weight = torch.randn(128, dtype=torch.bfloat16)
+    z = torch.randn_like(x)
+
+    x_fp32 = x.float()
+    rstd_ref = torch.rsqrt((x_fp32 * x_fp32).mean(dim=-1) + eps)
+    y_ref = (x_fp32 * rstd_ref[:, None] * weight.float()) * _silu(z.float())
+
+    y_npu, _, rstd_npu = custom_ops.layer_norm_fwd_npu(
+        x.npu(), weight.npu(), None, z.npu(), eps, -1, True, True)
+
+    assert torch.allclose(y_npu.cpu().float(), y_ref, atol=1e-2, rtol=1e-2)
+    assert torch.allclose(rstd_npu.cpu().float(), rstd_ref, atol=1e-2,
+                          rtol=1e-2)

xllm_ops/build_aclnn.sh

@@ -34,8 +34,11 @@ get_cann_toolkit_version() {
     for version_file in \
         "${ASCEND_TOOLKIT_HOME}/toolkit/version.info" \
         "${ASCEND_TOOLKIT_HOME}/version.info" \
+        "${ASCEND_TOOLKIT_HOME}/compiler/version.info" \
         "${ASCEND_HOME_PATH}/toolkit/version.info" \
-        "${ASCEND_HOME_PATH}/version.info"; do
+        "${ASCEND_HOME_PATH}/version.info" \
+        "${ASCEND_HOME_PATH}/opp/version.info" \
+        "${ASCEND_HOME_PATH}/compiler/version.info"; do
         if [[ -f "${version_file}" ]]; then
             version_line=$(grep -m1 '^Version=' "${version_file}" 2>/dev/null || true)
             if [[ -n "${version_line}" ]]; then

xllm_ops/layer_norm_fwd/op_kernel/layer_norm_fwd.cpp

@@ -26,6 +26,11 @@ constexpr uint32_t kQwenGroupSize = 128;
 constexpr uint32_t kN128ReduceBatchRows = 32;
 constexpr float kInvQwenGroupSize = 0.0078125f;
 constexpr float kDefaultEps = 1e-6f;
+#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510
+constexpr bool kNeedsExplicitScalarSync = true;
+#else
+constexpr bool kNeedsExplicitScalarSync = false;
+#endif
 
 enum KernelMode : uint32_t {
   kFullRow = 0,
@@ -52,44 +57,60 @@ __aicore__ inline float NormalizeEps(float eps) {
   return eps > 0.0f ? eps : kDefaultEps;
 }
 
+template <HardEvent event>
+__aicore__ inline event_t ResolveEventId(event_t fixed_event_id) {
+#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510
+  return fixed_event_id;
+#else
+  return static_cast<event_t>(GetTPipePtr()->FetchEventID(event));
+#endif
+}
+
 __aicore__ inline void WaitMte2ToV() {
-  event_t event_id = static_cast<event_t>(EVENT_ID0);
+  const event_t event_id =
+      ResolveEventId<HardEvent::MTE2_V>(static_cast<event_t>(EVENT_ID0));
   SetFlag<HardEvent::MTE2_V>(event_id);
   WaitFlag<HardEvent::MTE2_V>(event_id);
 }
 
 __aicore__ inline void WaitSToMte3() {
-  event_t event_id = static_cast<event_t>(EVENT_ID1);
+  const event_t event_id =
+      ResolveEventId<HardEvent::S_MTE3>(static_cast<event_t>(EVENT_ID1));
   SetFlag<HardEvent::S_MTE3>(event_id);
   WaitFlag<HardEvent::S_MTE3>(event_id);
 }
 
 __aicore__ inline void WaitVToS() {
-  event_t event_id = static_cast<event_t>(EVENT_ID7);
+  const event_t event_id =
+      ResolveEventId<HardEvent::V_S>(static_cast<event_t>(EVENT_ID7));
   SetFlag<HardEvent::V_S>(event_id);
   WaitFlag<HardEvent::V_S>(event_id);
 }
 
 __aicore__ inline void WaitSToV() {
-  event_t event_id = static_cast<event_t>(EVENT_ID6);
+  const event_t event_id =
+      ResolveEventId<HardEvent::S_V>(static_cast<event_t>(EVENT_ID6));
   SetFlag<HardEvent::S_V>(event_id);
   WaitFlag<HardEvent::S_V>(event_id);
 }
 
 __aicore__ inline void WaitVToMte3() {
-  event_t event_id = static_cast<event_t>(EVENT_ID2);
+  const event_t event_id =
+      ResolveEventId<HardEvent::V_MTE3>(static_cast<event_t>(EVENT_ID2));
   SetFlag<HardEvent::V_MTE3>(event_id);
   WaitFlag<HardEvent::V_MTE3>(event_id);
 }
 
 __aicore__ inline void WaitVToMte2() {
-  event_t event_id = static_cast<event_t>(EVENT_ID3);
+  const event_t event_id =
+      ResolveEventId<HardEvent::V_MTE2>(static_cast<event_t>(EVENT_ID3));
   SetFlag<HardEvent::V_MTE2>(event_id);
   WaitFlag<HardEvent::V_MTE2>(event_id);
 }
 
 __aicore__ inline void WaitMte3ToV() {
-  event_t event_id = static_cast<event_t>(EVENT_ID4);
+  const event_t event_id =
+      ResolveEventId<HardEvent::MTE3_V>(static_cast<event_t>(EVENT_ID4));
   SetFlag<HardEvent::MTE3_V>(event_id);
   WaitFlag<HardEvent::MTE3_V>(event_id);
 }
@@ -173,29 +194,44 @@ __aicore__ inline float ReduceSumValue(LocalTensor<float> scalar,
   PipeBarrier<PIPE_V>();
   ReduceSum(scalar, backup, reduce_tmp, count);
   PipeBarrier<PIPE_V>();
-  WaitVToS();
+  if constexpr (kNeedsExplicitScalarSync) {
+    WaitVToS();
+  }
   return scalar.GetValue(0);
 }
 
 __aicore__ inline float ReduceSumValueNoCopy(LocalTensor<float> scalar,
                                              LocalTensor<float> src,
                                              LocalTensor<float> reduce_tmp,
                                              uint32_t count) {
-  WaitSToV();
+  if constexpr (kNeedsExplicitScalarSync) {
+    WaitSToV();
+  }
   ReduceSum(scalar, src, reduce_tmp, count);
   PipeBarrier<PIPE_V>();
-  WaitVToS();
-  float v = scalar.GetValue(0);
-  WaitSToV();
-  return v;
+  if constexpr (kNeedsExplicitScalarSync) {
+    WaitVToS();
+    const float value = scalar.GetValue(0);
+    WaitSToV();
+    return value;
+  }
+  return scalar.GetValue(0);
 }
 
 __aicore__ inline float SqrtValue(LocalTensor<float> scalar, float value) {
-  WaitVToS();
+  if constexpr (kNeedsExplicitScalarSync) {
+    WaitVToS();
+  }
   scalar.SetValue(0, value);
-  WaitSToV();
+  if constexpr (kNeedsExplicitScalarSync) {
+    WaitSToV();
+  }
   Sqrt(scalar, scalar, 1);
-  WaitVToS();
+  if constexpr (kNeedsExplicitScalarSync) {
+    WaitVToS();
+  } else {
+    PipeBarrier<PIPE_V>();
+  }
   return scalar.GetValue(0);
 }
 
@@ -437,7 +473,9 @@ class KernelFullRow {
     }
     pipe_.InitBuffer(xFp32Buf_, tile_elems * sizeof(float));
     pipe_.InitBuffer(tmpFp32Buf_, tile_elems * sizeof(float));
-    pipe_.InitBuffer(reduceTmpBuf_, tileRows_ * AlignUp(groupAlign_, kFp32BlockElems) * sizeof(float));
+    // Rows are processed serially, so they can share one reduction scratch.
+    pipe_.InitBuffer(reduceTmpBuf_,
+                     AlignUp(groupAlign_, kFp32BlockElems) * sizeof(float));
     if (IsRmsGateN128() || IsRmsN128NoZ()) {
       pipe_.InitBuffer(rowReduceTmpBuf_,
                        kN128ReduceBatchRows * kFp32RepeatElems * sizeof(float));
@@ -804,7 +842,7 @@ class KernelFullRow {
     for (uint32_t r = 0; r < rows; ++r) {
       LocalTensor<float> row_x = x_fp32[r * groupAlign_];
       LocalTensor<float> row_tmp = tmp_fp32[r * groupAlign_];
-      LocalTensor<float> row_reduce = reduce_tmp[r * AlignUp(groupAlign_, kFp32BlockElems)];
+      LocalTensor<float> row_reduce = reduce_tmp;
 
       if (hasZ_ != 0 && normBeforeGate_ == 0) {
         CastToFloat<T>(row_tmp, z_local[r * groupAlign_], groupAlign_);