Skip to content

[KMCompiler][ttx] Optimize gelu with rowwise nomask kernels#364

Open
YangLong114514 wants to merge 2 commits into
XPU-Forces:masterfrom
YangLong114514:KMCompiler-GeLU
Open

[KMCompiler][ttx] Optimize gelu with rowwise nomask kernels#364
YangLong114514 wants to merge 2 commits into
XPU-Forces:masterfrom
YangLong114514:KMCompiler-GeLU

Conversation

@YangLong114514

@YangLong114514 YangLong114514 commented Jun 15, 2026

Copy link
Copy Markdown

Description

The gelu operator has been optimized for the Ascend platform.

Changes

  1. Updated the BLOCK_SIZE_N calculation to use a capped power-of-two tile size. For GELU, the cap is set to 1024 to reduce UB pressure from the tanh approximation path.

  2. Added dynamic grid calculation based on the autotuned BLOCK_SIZE_M, reducing empty programs for small shapes.

  3. Added no-mask forward kernels for fully divisible row/column tiles, removing mask construction and masked tl.load / tl.store overhead.

  4. Added a no-mask single-tile forward kernel for shapes where n_cols == BLOCK_SIZE_N, eliminating the inner column loop for common benchmark shapes.

Performance

Using Ascend 910C and Triton 3.2.x of FlagTree:

shape before after speedup = before/after
(16, 16) 2.560 1.512 1.69
(32, 64) 2.728 1.516 1.80
(128, 128) 3.324 2.036 1.63
(128, 256) 3.7602 2.400 1.57
(256, 256) 4.4282 3.304 1.34
(512, 256) 4.984 4.084 1.22

Accuracy test

The correctness verification of the GeLU operator for all three data types, float32, float16, and bf16, passed in the Mojo correctness verification directory.
image

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Code Review

This pull request optimizes the GELU forward pass on NPU by introducing specialized no-mask kernels (_gelu_fwd_nomask_kernel and _gelu_fwd_nomask_single_kernel) when dimensions align, and refactors grid and block size calculations. The feedback suggests reusing the cached get_num_cores utility from .utils instead of dynamically querying device properties on every grid calculation, and dynamically computing GELU_TANH_MAX_BLOCK_SIZE_M from the autotune configurations to prevent potential out-of-bounds bugs if the configurations are updated in the future.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread mojo_opset/backends/ttx/kernels/npu/gelu.py Outdated
Comment thread mojo_opset/backends/ttx/kernels/npu/gelu.py Outdated
Comment thread mojo_opset/backends/ttx/kernels/npu/gelu.py Outdated
@YangLong114514 YangLong114514 changed the title [KMCompiler] Optimize gelu with rowwise nomask kernels [KMCompiler][ttx] Optimize gelu with rowwise nomask kernels Jun 15, 2026
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