[Model] Add Afmoe architecture implementation #28332

pranav4501 · 2025-11-08T00:11:04Z

Purpose

This PR adds architecture implementation of upcoming Arcee AI AFMoE (trinity) models.

Test Plan

The model is not public yet, verified serving of AFMoE across configs.

Test Result

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
(Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Signed-off-by: Maziyar Panahi <maziyar.panahi@iscpif.fr>

github-actions · 2025-11-08T00:11:11Z

👋 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 ask your reviewers to trigger select CI tests on top of fastcheck CI.

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.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

mergify · 2025-11-08T00:11:45Z

Documentation preview: https://vllm--28332.org.readthedocs.build/en/28332/

Signed-off-by: Pranav <veldurthipranav@gmail.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

vllm/model_executor/models/afmoe.py

Signed-off-by: Pranav <veldurthipranav@gmail.com>

heheda12345 · 2025-11-11T08:06:51Z

Also CC @ywang96

mgoin · 2025-11-13T23:19:36Z

vllm/model_executor/models/afmoe.py

+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        if self.n_shared_experts > 0:
+            shared_output = self.shared_experts(hidden_states)
+        else:
+            shared_output = torch.zeros_like(hidden_states)
+
+        router_logits = self.gate(hidden_states.to(dtype=torch.float32))
+
+        fused_moe_out = self.experts(
+            hidden_states=hidden_states, router_logits=router_logits
+        )
+
+        if self.shared_experts is not None:
+            shared_output, final_hidden_states = fused_moe_out
+            final_hidden_states = final_hidden_states + shared_output
+        else:
+            final_hidden_states = fused_moe_out
+        if self.tp_size > 1:
+            final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
+                final_hidden_states
+            )
+
+        return final_hidden_states.view(num_tokens, hidden_dim)


I think you still need to update this to work with SharedFusedMoE. See qwen2_moe for instance

vllm/vllm/model_executor/models/qwen2_moe.py

Lines 170 to 188 in e64011f

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

# NOTE: hidden_states can have either 1D or 2D shape.

orig_shape = hidden_states.shape

hidden_dim = hidden_states.shape[-1]

hidden_states = hidden_states.view(-1, hidden_dim)

# router_logits: (num_tokens, n_experts)

router_logits, _ = self.gate(hidden_states)

final_hidden_states = self.experts(

hidden_states=hidden_states, router_logits=router_logits

)

if self.shared_expert is not None:

final_hidden_states = final_hidden_states[0] + final_hidden_states[1]

if self.tp_size > 1:

final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501

final_hidden_states

)

return final_hidden_states.view(orig_shape)

Yeah, I missed removing the previous shared experts computation.
Is that all or am i missing something else?

cc @bnellnm @alexm-redhat could you check too?

Looks good to me.

mgoin · 2025-11-13T23:23:02Z

Also should make sure to add this arch to tests/models/registry.py with is_available_online=False so we can easily add it to the basic models test when ready

DarkLight1337 · 2025-11-14T05:55:52Z

vllm/model_executor/models/afmoe.py

+                logical_replica_count=logical_replica_count,
+            )
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:


This was renamed to embed_input_ids in #27583, can you update this?

Signed-off-by: Pranav <veldurthipranav@gmail.com>

pranav4501 · 2025-11-14T21:57:19Z

Made the fixes
DP seems to be broken since merging latest main
patching in #28740 seems to fix it

mgoin

LGTM and ready to go! I think we can merge ahead of time and enable the model test later

Signed-off-by: Maziyar Panahi <maziyar.panahi@iscpif.fr> Signed-off-by: Pranav <veldurthipranav@gmail.com> Co-authored-by: Maziyar Panahi <maziyar.panahi@iscpif.fr>

Signed-off-by: Maziyar Panahi <maziyar.panahi@iscpif.fr> Signed-off-by: Pranav <veldurthipranav@gmail.com> Co-authored-by: Maziyar Panahi <maziyar.panahi@iscpif.fr> Signed-off-by: jiang1.li <jiang1.li@intel.com>

Signed-off-by: Maziyar Panahi <maziyar.panahi@iscpif.fr> Signed-off-by: Pranav <veldurthipranav@gmail.com> Co-authored-by: Maziyar Panahi <maziyar.panahi@iscpif.fr>

afmoe modeling and config

9b859c9

Signed-off-by: Maziyar Panahi <maziyar.panahi@iscpif.fr>

mergify bot added documentation Improvements or additions to documentation new-model Requests to new models labels Nov 8, 2025

add docs

f8ecf32

Signed-off-by: Pranav <veldurthipranav@gmail.com>

pranav4501 force-pushed the afmoe-modeling branch from 914ad26 to f8ecf32 Compare November 8, 2025 00:23

pranav4501 marked this pull request as ready for review November 8, 2025 00:29

chatgpt-codex-connector bot reviewed Nov 8, 2025

View reviewed changes

vllm/model_executor/models/afmoe.py Show resolved Hide resolved

mgoin reviewed Nov 9, 2025

View reviewed changes

vllm/model_executor/models/afmoe.py Outdated Show resolved Hide resolved

Update MoE to SharedFusedMoE

48520e4

Signed-off-by: Pranav <veldurthipranav@gmail.com>

pranav4501 requested a review from mgoin November 10, 2025 19:51

mgoin reviewed Nov 13, 2025

View reviewed changes

mgoin requested review from DarkLight1337 and ywang96 November 13, 2025 23:21

DarkLight1337 reviewed Nov 14, 2025

View reviewed changes

pranav4501 added 4 commits November 14, 2025 12:06

Merge branch 'main' into afmoe-modeling

1b3986f

add model to test registry

5fe6ce3

Signed-off-by: Pranav <veldurthipranav@gmail.com>

update embed refactor

e57e2b9

Signed-off-by: Pranav <veldurthipranav@gmail.com>

remove redundant shared experts

db7559d

Signed-off-by: Pranav <veldurthipranav@gmail.com>

mgoin approved these changes Nov 17, 2025

View reviewed changes

mgoin added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 17, 2025

Merge branch 'main' into afmoe-modeling

5764996

vllm-bot merged commit f77bce0 into vllm-project:main Nov 17, 2025
48 of 50 checks passed

	def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	# NOTE: hidden_states can have either 1D or 2D shape.
	orig_shape = hidden_states.shape
	hidden_dim = hidden_states.shape[-1]
	hidden_states = hidden_states.view(-1, hidden_dim)

	# router_logits: (num_tokens, n_experts)
	router_logits, _ = self.gate(hidden_states)
	final_hidden_states = self.experts(
	hidden_states=hidden_states, router_logits=router_logits
	)
	if self.shared_expert is not None:
	final_hidden_states = final_hidden_states[0] + final_hidden_states[1]
	if self.tp_size > 1:
	final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501
	final_hidden_states
	)

	return final_hidden_states.view(orig_shape)

Uh oh!

[Model] Add Afmoe architecture implementation #28332

[Model] Add Afmoe architecture implementation #28332

Uh oh!

Conversation

pranav4501 commented Nov 8, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

github-actions bot commented Nov 8, 2025

Uh oh!

mergify bot commented Nov 8, 2025

Uh oh!

chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

heheda12345 commented Nov 11, 2025

Uh oh!

mgoin Nov 13, 2025

Choose a reason for hiding this comment

Uh oh!

pranav4501 Nov 14, 2025

Choose a reason for hiding this comment

Uh oh!

mgoin Nov 14, 2025

Choose a reason for hiding this comment

Uh oh!

bnellnm Nov 14, 2025

Choose a reason for hiding this comment

Uh oh!

mgoin commented Nov 13, 2025

Uh oh!

DarkLight1337 Nov 14, 2025

Choose a reason for hiding this comment

Uh oh!

pranav4501 commented Nov 14, 2025

Uh oh!

mgoin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

pranav4501 commented Nov 8, 2025 •

edited by github-actions bot

Loading