Add dropping strategy

[RissyRan]

Description

Enable token dropping for matmul implementation

Next steps:

• work on dispatch & combine mask - make expert parallelism easy
• add load balance loss

Test

Have a unite test to check a single layer output (become dropless if capacity factor is large enough):

• capacity factor = 2 as default (it should be 1 ideally with evenly distributed) - link - pass
• capacity factor = 0.1 (extreme case) - link - fail

[ZhiyuLi-goog]

Thank you @RissyRan for adding dropping strategy!! I just added some nit.

[RissyRan]

Thank you @RissyRan for adding dropping strategy!! I just added some nit.

Thanks Zhiyu! Have you published comments?

[ZhiyuLi-goog]

Thank you @RissyRan for adding dropping strategy!! I just added some nit.

Thanks Zhiyu! Have you published comments?

Just published.