feat(ascend): add SwiGLU forward and backward kernels - #381
Conversation
Signed-off-by: chenyang <2082464740@qq.com>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Team Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
zhangj1an
left a comment
There was a problem hiding this comment.
Thank you for your contribution! The gtest contract passes for fp32/bf16/fp16 (forward + backward) and the pytest suite passes.
Gtest results (scripts/check_operator.py, all three dtypes, --check-grad):
| dtype | forward | gradient:gate | gradient:up | verdict |
|---|---|---|---|---|
| fp32 | 0.0 (bitwise) | 9.54e-7 (~1 ulp) | 0.0 (bitwise) | ✅ PASS |
| bf16 | 3.01e-2 | 6.10e-5 | 0.0 (bitwise) | ✅ PASS |
| fp16 | 3.83e-3 | 4.88e-4 | 0.0 (bitwise) | ✅ PASS |
Thresholds are (atol, rtol) = (2e-2, 1.6e-2) for bf16 and (1e-3, 1e-3) for fp16, with per-element allclose semantics (|diff| ≤ atol + rtol·|ref|). So bf16 forward max_abs 3.01e-2 and fp16 forward max_abs 3.83e-3 exceed the bare atol but pass legitimately via the rtol term — that is the intended gtest contract behavior, not a fudge.
Command used:
ASCEND_RT_VISIBLE_DEVICES=1 python scripts/check_operator.py \
--op swiglu --candidate ascend --device npu \
--dtype {fp32,bf16,fp16} --batch 2 --seq 16 --vocab 257 --normalized-dim 4096 \
--check-gradOptional minor Issues
- Flaky test (must fix):
test_autograd_wrapper_contiguity_and_gradient_routingfails in ~2/5 full-suite runs (different parametrization each time, always passes in isolation):torch.testing.assert_close(result, ref, rtol=0, atol=0)with 1/35 elements off by 1 fp32 ULP (5.96e-8). The wrapper path feeds contiguous copies toNativeSwiGLUOpwhile the reference runs on the original strided.t()tensors; torch's elementwise kernels can differ by 1 ULP between layouts. CPU-only wrapper test — fix by feeding the same layout to both paths or allowing 1 ULP. - Naming inconsistent with #378: the binding file is
csrc/ascend/bindings.ascwhile #378 usescsrc/ascend/ops_ascend.cpp(other Ascend branches usenpu_module.cpp) — three names for the same role; also a host-only pybind TU named.ascrides through the bisheng kernel glob. The bound symbolsswiglu_forward/swiglu_backwardlack the_ascendsuffix used byrope_apply_ascendandbatch_invariant_logp_ascendin the same_C_npumodule.
Summary
gate * sigmoid(gate) * up) forward and backward kernels written in Ascend C (csrc/ascend/activation.asc), computing in FP32 and returning the input dtype (fp16 / bf16 / fp32)csrc/ascend/bindings.ascmodule initializer (swiglu_forward/swiglu_backward) with.pyistubsSwiGLUAscendOpPython wrapper with autograd support (once_differentiable), input validation, arbitrary shapes / empty tensors / strided-view handling, and an FP32-output variantASCEND_SWIGLUintoKernelRegistryNPU dispatch priority (ascend -> pytorch fallback) and update gtest operator specstests/test_swiglu.py(precision vs reference, dtype/shape/strided/empty coverage, autograd checks) and updatedocs/operators/activation.md