[Model]: Fused MoE for nomic-embed-text-v2-moe

[Isotr0py]

• Following PR for [New Model]: nomic-embed-text-v2-moe #17785
• Add ungated activation support to fused_moe.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[DarkLight1337]

cc @noooop can you help test this?

[noooop]

Sorry, I'm going outside currently and can only go back for testing on the 27th

[DarkLight1337]

Since this model already has a test in CI, I'll just unblock it and see if it passes

[noooop]

https://huggingface.co/Alibaba-NLP/gte-Qwen2-1.5B-instruct/discussions/37

Fix: Respect is_causal=False config in forward to enable bidirectional attention

The test failure is related to this, and after a year, it has finally been fixed.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Isotr0py.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[noooop]

sorry for late response

language-models-test-extended has verified the model's results on mteb/STS12

I tested a larger mteb/T2Reranking dataset.
There can still be significant speed improvements without kernel tuning.

Supporting long context requires #18755

# https://github.com/noooop/snippet/blob/main/benchmarks/test_mteb/test_speed.py

python -m benchmarks.test_mteb.test_speed nomic-ai/nomic-embed-text-v2-moe float16 512
python -m benchmarks.test_mteb.test_speed nomic-ai/nomic-embed-text-v2-moe float32 512

# https://github.com/noooop/snippet/blob/main/benchmarks/test_mteb/test_speed_long.py

python -m benchmarks.test_mteb.test_speed_long nomic-ai/nomic-embed-text-v2-moe float16
python -m benchmarks.test_mteb.test_speed_long nomic-ai/nomic-embed-text-v2-moe float32

main:
512
float32 0.6708810434272698 480.56it/s
float16 0.7195732012356584 813.04it/s

8192
float32 0.6611511744670271 186.95it/s
float16 0.7074889328412163 396.59it/s

this pr:
512
float32 0.6709708587757305 610.92it/s 127% ↑
float16 0.7196080045872064 1234.22it/s 151% ↑

8192
float32 0.6610671251625739 220.47it/s 118% ↑
float16 0.7075334132753368 580.60it/s 146% ↑

LGTM

[Isotr0py]

Personally, is_gated is better than is_act_and_mul

Yes, I have considered is_gated at first. But since we're using "gating" in gating_output, I'm afraid that it will cause some misunderstanding for those who are new to MoE, so I changed it to is_act_and_mul which is more intuitive.

[DarkLight1337]

Any update on this? It would help with getting the model to support V1 with torch.compile

[Isotr0py]

Let me update this PR to catch recent MoE refactoring.

[Isotr0py]

I have confirmed tests/models/language/pooling/test_nomic.py can still pass with this update.

[DarkLight1337]

Thanks for the quick update!

[noooop]

buildkite/ci/pr/language-models-test-extended-pooling

Failed tests can be fixed by #21747