Skip to content

[NPU, A3] Add NPU kernel support for A3 machines [1/2] - #1220

Merged
Tcc0403 merged 1 commit into
linkedin:mainfrom
pt-ecosystem:main
May 13, 2026
Merged

[NPU, A3] Add NPU kernel support for A3 machines [1/2]#1220
Tcc0403 merged 1 commit into
linkedin:mainfrom
pt-ecosystem:main

Conversation

@zheliuyu

@zheliuyu zheliuyu commented May 11, 2026

Copy link
Copy Markdown
Contributor

Motivation

This work follows the roadmap in linkedin/Liger-Kernel#969. The goal is to exercise the NPU kernel on Atlas 800T A3 (64G) and report how the test suite behaves on that hardware.

Details

  • Fixed the Ascend implementation of attn_res so it passes on A3 alongside the rest of the suite.
  • After the fix, the full test run completes successfully. 🍾
image

Why attn_res failed on A3

The failure showed up as vector-core / ACL errors (e.g. 507035, device sync failing), not a normal atol/rtol mismatch.

  • Ascend attn_res uses wide masked loads along the feature dim (e.g. BLOCK_D = next_power_of_2(D)).
  • Tests include awkward sizes like D = 123 with float32, so row pitch is 123 × 4 = 492 bytes—not 32B/64B-friendly for many vectorized paths on this stack, which can trigger vector-core faults for that lowering.
  • Fix: pad the last dim to a multiple of 16, pass d_stride as the real memory pitch, keep D for math/masks, and slice/pad tensors so callers still see logical D.

Benchmark results for the 4 most frequently used kernels

cross_entropy_memory_full_token_length cross_entropy_speed_backward_token_length
cross_entropy_memory_full_token_length cross_entropy_speed_backward_token_length
cross_entropy_speed_forward_token_length cross_entropy_speed_full_token_length
cross_entropy_speed_forward_token_length cross_entropy_speed_full_token_length
cross_entropy_speed_no-grad-forward_token_length rms_norm_memory_full_token_length
cross_entropy_speed_no-grad-forward_token_length rms_norm_memory_full_token_length
rms_norm_speed_backward_token_length rms_norm_speed_forward_token_length
rms_norm_speed_backward_token_length rms_norm_speed_forward_token_length
rms_norm_speed_full_token_length rope_memory_full_token_length
rms_norm_speed_full_token_length rope_memory_full_token_length
rope_speed_backward_token_length rope_speed_forward_token_length
rope_speed_backward_token_length rope_speed_forward_token_length
rope_speed_full_token_length swiglu_memory_full_token_length
rope_speed_full_token_length swiglu_memory_full_token_length
swiglu_speed_backward_token_length swiglu_speed_forward_token_length
swiglu_speed_backward_token_length swiglu_speed_forward_token_length
swiglu_speed_full_token_length
swiglu_speed_full_token_length

Testing Done

  • Hardware Type: All NPUs.
  • run make test to ensure correctness
  • run make checkstyle to ensure code style
  • run make test-convergence to ensure convergence

@Tcc0403 Tcc0403 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.

LGTM

@zheliuyu
zheliuyu marked this pull request as ready for review May 13, 2026 12:38
@zheliuyu
zheliuyu requested a review from Tcc0403 May 13, 2026 12:38
@zheliuyu

Copy link
Copy Markdown
Contributor Author

@Tcc0403 This PR is ready for review.

@Tcc0403 Tcc0403 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.

LGTM

@Tcc0403
Tcc0403 added this pull request to the merge queue May 13, 2026
Merged via the queue into linkedin:main with commit 38c0d4f May 13, 2026
5 of 7 checks passed
@zheliuyu zheliuyu changed the title [NPU, A3] Add NPU kernel support for A3 machines [NPU, A3] Add NPU kernel support for A3 machines [1/2] May 18, 2026
@zheliuyu zheliuyu mentioned this pull request May 25, 2026
3 tasks
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.

2 participants