Skip to content

use hwy/simd for RMSNorm(f, bf, f) calculation#78

Merged
copybara-service[bot] merged 2 commits intogoogle:devfrom
enum-class:rmsnorm2
Mar 10, 2024
Merged

use hwy/simd for RMSNorm(f, bf, f) calculation#78
copybara-service[bot] merged 2 commits intogoogle:devfrom
enum-class:rmsnorm2

Conversation

@enum-class
Copy link
Contributor

use highway/simd for RMSNorm(f, bf, f) calculation for performance purpose.

static HWY_NOINLINE HWY_MAYBE_UNUSED void RMSNorm(
     const float* HWY_RESTRICT x, const hwy::bfloat16_t* HWY_RESTRICT weight,
     float* HWY_RESTRICT out, size_t size) ;

Copy link
Member

@jan-wassenberg jan-wassenberg left a comment

Choose a reason for hiding this comment

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

Very nice, thanks for implementing! One style suggestion:

ops.h Outdated
for (size_t j = 0; j < size; j++) {
// Note 1.0f centering here
out[j] = (1.0f + hwy::F32FromBF16(weight[j])) * (ss * x[j]);
constexpr size_t unroll_size = 2;
Copy link
Member

Choose a reason for hiding this comment

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

For constexpr please use kUnrollSize naming :)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I 've set up Clang-Tidy to enforce the kCamelCase style for constexpr variables in the codebase, aiming for automated style checking and corrections.
Maybe in future it can be applied to fix all constexpr styles

Copy link
Member

@jan-wassenberg jan-wassenberg left a comment

Choose a reason for hiding this comment

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

Nice, great that you have updated the lint to check for this :)

@jan-wassenberg jan-wassenberg added the copybara-import Trigger Copybara for merging pull requests label Mar 6, 2024
@enum-class
Copy link
Contributor Author

It seems copybara has faced an issue, maybe the rerun can solve it :))

@jan-wassenberg
Copy link
Member

Yes indeed, still have an issue with copybara transforms. Looking into it :)

@jan-wassenberg jan-wassenberg added copybara-import Trigger Copybara for merging pull requests and removed copybara-import Trigger Copybara for merging pull requests labels Mar 7, 2024
@jan-wassenberg
Copy link
Member

Issue resolved 👍 One more internal review, then it will be merged.

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