Skip to content

[KernelGen][Ascend] Add gemma_rms_norm operator with Triton kernel for Ascend#14

Open
yzw1128 wants to merge 1 commit into
flagos-ai:masterfrom
yzw1128:feature/ascend-gemma-rms-norm
Open

[KernelGen][Ascend] Add gemma_rms_norm operator with Triton kernel for Ascend#14
yzw1128 wants to merge 1 commit into
flagos-ai:masterfrom
yzw1128:feature/ascend-gemma-rms-norm

Conversation

@yzw1128

@yzw1128 yzw1128 commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

PR Category

Operator

Type of Change

New Feature

Description

Add Ascend-optimized Gemma RMSNorm Triton kernel implementation under the _ascend backend ops directory.

The implementation provides:

  • _gemma_rmsnorm_kernel: standard RMSNorm with (1 + w) scaling
  • _gemma_fused_add_rmsnorm_kernel: fused residual-add + RMSNorm
  • gemma_rms_norm(x, weight, eps, residual): unified entry point matching the existing tensor-based API

Issue

N/A

Progress

  • Change is properly reviewed (1 reviewer required, 2 recommended).
  • Change is responded to an issue.
  • Change is fully covered by a UT.

Performance

Ascend NPU optimized implementation.

🤖 Generated with Claude Code

…r Ascend

Add Ascend-optimized Gemma RMSNorm Triton kernel implementation
under the _ascend backend ops directory.

The implementation provides:
- _gemma_rmsnorm_kernel: standard RMSNorm with (1 + w) scaling
- _gemma_fused_add_rmsnorm_kernel: fused residual-add + RMSNorm
- gemma_rms_norm: unified entry point matching the existing tensor-based
  API (x, weight, eps, residual)
@yzw1128 yzw1128 force-pushed the feature/ascend-gemma-rms-norm branch from f122f34 to 63d63e3 Compare July 8, 2026 03:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant