Skip to content

Add NEON fast path for SFP weight decode - #971

Open
RaviTriv wants to merge 4 commits into
google:devfrom
RaviTriv:neon-weight-decode
Open

Add NEON fast path for SFP weight decode#971
RaviTriv wants to merge 4 commits into
google:devfrom
RaviTriv:neon-weight-decode

Conversation

@RaviTriv

Copy link
Copy Markdown

Proposed fast path for weight decode on NEON that uses table look up and shift. About the same ops but looks like we can sneak in a few cheaper ones, roughly 1.1x faster on SFP matmuls.

@RaviTriv
RaviTriv marked this pull request as ready for review August 4, 2026 00:25
jan-wassenberg
jan-wassenberg previously approved these changes Aug 5, 2026

@jan-wassenberg jan-wassenberg left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Nice work and notable speedup, thanks for adding :) Have you measured on non-AVX3_DL x86 to see if your approach is faster than the generic? If so, we could replace our generic fallback on all platforms, rather than just NEON. SVE is also increasingly available/chosen.

@jan-wassenberg jan-wassenberg added the copybara-import Trigger Copybara for merging pull requests label Aug 5, 2026
@RaviTriv

RaviTriv commented Aug 6, 2026

Copy link
Copy Markdown
Author

non-AVX3_DL x86

thanks!

good idea, have not measured. taking a look :)

@RaviTriv

RaviTriv commented Aug 7, 2026

Copy link
Copy Markdown
Author

So looks like this is also a good speedup on SVE, roughly 1.1x. The default fallback is still better for AVX2 and AVX3 because x86 doesnt have a per-lane variable shifts for u8.

@jan-wassenberg jan-wassenberg left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Nice, makes sense SVE2_128 also works well. Would also think SVE2 works, any reason we limit it to just SVE2_128?

@jan-wassenberg jan-wassenberg added copybara-import Trigger Copybara for merging pull requests and removed copybara-import Trigger Copybara for merging pull requests labels Aug 7, 2026
@RaviTriv

RaviTriv commented Aug 7, 2026

Copy link
Copy Markdown
Author

Yep, only enabled for SVE2_128 because I was referring to comment indicating that SVE2_128 is the only enabled SVE target right now?

@jan-wassenberg

Copy link
Copy Markdown
Member

True, though this is due to compiler bugs which will hopefully be fixed at some point. I think there is no need to disable/exclude those SVE targets here as well, but let's still land as-is.

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

Labels

copybara-import Trigger Copybara for merging pull requests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants