Skip to content

feat(ascend): add SwiGLU forward and backward kernels - #381

Open
erfgss wants to merge 1 commit into
RL-Align:testfrom
erfgss:feat/swiGLU_ascend
Open

feat(ascend): add SwiGLU forward and backward kernels#381
erfgss wants to merge 1 commit into
RL-Align:testfrom
erfgss:feat/swiGLU_ascend

Conversation

@erfgss

@erfgss erfgss commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Add fused SwiGLU (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)
  • Register the kernels on the NPU extension via a new csrc/ascend/bindings.asc module initializer (swiglu_forward / swiglu_backward) with .pyi stubs
  • Add SwiGLUAscendOp Python wrapper with autograd support (once_differentiable), input validation, arbitrary shapes / empty tensors / strided-view handling, and an FP32-output variant
  • Wire ASCEND_SWIGLU into KernelRegistry NPU dispatch priority (ascend -> pytorch fallback) and update gtest operator specs
  • Extend tests/test_swiglu.py (precision vs reference, dtype/shape/strided/empty coverage, autograd checks) and update docs/operators/activation.md

Signed-off-by: chenyang <2082464740@qq.com>
@coderabbitai

coderabbitai Bot commented Sep 3, 2026

Copy link
Copy Markdown

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Team

Run ID: 1bd775d9-4a75-4199-b759-afbc8f2ab30f

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

Comment @coderabbitai help to get the list of available commands.

@zhangj1an zhangj1an left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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-grad

Optional minor Issues

  1. Flaky test (must fix): test_autograd_wrapper_contiguity_and_gradient_routing fails 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 to NativeSwiGLUOp while 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.
  2. Naming inconsistent with #378: the binding file is csrc/ascend/bindings.asc while #378 uses csrc/ascend/ops_ascend.cpp (other Ascend branches use npu_module.cpp) — three names for the same role; also a host-only pybind TU named .asc rides through the bisheng kernel glob. The bound symbols swiglu_forward / swiglu_backward lack the _ascend suffix used by rope_apply_ascend and batch_invariant_logp_ascend in the same _C_npu module.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants