diff --git a/test/python_test/test_layer_norm_fwd.py b/test/python_test/test_layer_norm_fwd.py index eb68a6b..e5cb245 100644 --- a/test/python_test/test_layer_norm_fwd.py +++ b/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) \ No newline at end of file + 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) diff --git a/xllm_ops/build_aclnn.sh b/xllm_ops/build_aclnn.sh index ae48ac3..30cddda 100644 --- a/xllm_ops/build_aclnn.sh +++ b/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 diff --git a/xllm_ops/layer_norm_fwd/op_kernel/layer_norm_fwd.cpp b/xllm_ops/layer_norm_fwd/op_kernel/layer_norm_fwd.cpp index 74dd783..3df1009 100644 --- a/xllm_ops/layer_norm_fwd/op_kernel/layer_norm_fwd.cpp +++ b/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 +__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(GetTPipePtr()->FetchEventID(event)); +#endif +} + __aicore__ inline void WaitMte2ToV() { - event_t event_id = static_cast(EVENT_ID0); + const event_t event_id = + ResolveEventId(static_cast(EVENT_ID0)); SetFlag(event_id); WaitFlag(event_id); } __aicore__ inline void WaitSToMte3() { - event_t event_id = static_cast(EVENT_ID1); + const event_t event_id = + ResolveEventId(static_cast(EVENT_ID1)); SetFlag(event_id); WaitFlag(event_id); } __aicore__ inline void WaitVToS() { - event_t event_id = static_cast(EVENT_ID7); + const event_t event_id = + ResolveEventId(static_cast(EVENT_ID7)); SetFlag(event_id); WaitFlag(event_id); } __aicore__ inline void WaitSToV() { - event_t event_id = static_cast(EVENT_ID6); + const event_t event_id = + ResolveEventId(static_cast(EVENT_ID6)); SetFlag(event_id); WaitFlag(event_id); } __aicore__ inline void WaitVToMte3() { - event_t event_id = static_cast(EVENT_ID2); + const event_t event_id = + ResolveEventId(static_cast(EVENT_ID2)); SetFlag(event_id); WaitFlag(event_id); } __aicore__ inline void WaitVToMte2() { - event_t event_id = static_cast(EVENT_ID3); + const event_t event_id = + ResolveEventId(static_cast(EVENT_ID3)); SetFlag(event_id); WaitFlag(event_id); } __aicore__ inline void WaitMte3ToV() { - event_t event_id = static_cast(EVENT_ID4); + const event_t event_id = + ResolveEventId(static_cast(EVENT_ID4)); SetFlag(event_id); WaitFlag(event_id); } @@ -173,7 +194,9 @@ __aicore__ inline float ReduceSumValue(LocalTensor scalar, PipeBarrier(); ReduceSum(scalar, backup, reduce_tmp, count); PipeBarrier(); - WaitVToS(); + if constexpr (kNeedsExplicitScalarSync) { + WaitVToS(); + } return scalar.GetValue(0); } @@ -181,21 +204,34 @@ __aicore__ inline float ReduceSumValueNoCopy(LocalTensor scalar, LocalTensor src, LocalTensor reduce_tmp, uint32_t count) { - WaitSToV(); + if constexpr (kNeedsExplicitScalarSync) { + WaitSToV(); + } ReduceSum(scalar, src, reduce_tmp, count); PipeBarrier(); - 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 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(); + } 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 row_x = x_fp32[r * groupAlign_]; LocalTensor row_tmp = tmp_fp32[r * groupAlign_]; - LocalTensor row_reduce = reduce_tmp[r * AlignUp(groupAlign_, kFp32BlockElems)]; + LocalTensor row_reduce = reduce_tmp; if (hasZ_ != 0 && normBeforeGate_ == 0) { CastToFloat(row_tmp, z_local[r * groupAlign_], groupAlign_);