Skip to content

feat(ascend): add SwiGLU operator - #334

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

feat(ascend): add SwiGLU operator#334
erfgss wants to merge 1 commit into
RL-Align:testfrom
erfgss:feat/ascend-swiglu

Conversation

@erfgss

@erfgss erfgss commented Aug 23, 2026

Copy link
Copy Markdown

Title:

feat(ascend): add SwiGLU operator

Summary

  • Add Ascend C forward and backward kernels for SwiGLU.
  • Support FP32, FP16, and BF16 inputs with FP32 intermediate computation.
  • Implement aligned multi-core 1D tiling, non-aligned tail handling, and empty tensor support.
  • Improve numerical stability with bounded sigmoid evaluation and reciprocal refinement.
  • Integrate the kernels with _C_npu, PyTorch Autograd, NPU dispatch, and OP_SPECS.
  • Add Ascend-specific tests and update the operator documentation.

Operator Definition

y      = silu(gate) * up
d_up   = dy * silu(gate)
d_gate = dy * up * sigmoid(gate)
         * (1 + gate * (1 - sigmoid(gate)))

gate and up must have the same shape, dtype, and device.

Supported Data Types

  • FP32
  • FP16
  • BF16

Activation and gradient calculations are performed in FP32, with results cast back to the input dtype.

Dispatch

SwiGLU dispatch priority on NPU:

Ascend C SwiGLU → PyTorch Native SwiGLU

Testing

73 passed, 95 skipped

The tests cover:

  • Forward and backward accuracy
  • FP32, FP16, and BF16 inputs
  • Non-aligned tail tiles
  • Batch and padding invariance
  • Non-contiguous tensors
  • Empty tensors
  • Shape and dtype validation
  • Registry and operator specification integration

@coderabbitai

coderabbitai Bot commented Aug 23, 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: Pro Plus

Run ID: a33f2fce-ea48-4893-9b1b-1a9e1fa36fa0

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

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

@zhangj1an

Copy link
Copy Markdown
Collaborator

thank you for your contribution! please fix ci conflicts. you can click on the "view details" tab shown below to see where went wrong
image

also, for every PR description, please paste pytest and gtest results to show your kernel is bitwise identical to pytorch reference.

thank you!

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