Skip to content

Commit c6c9830

Browse files
authored
[Bugfix] Mark 'hidden_states' as mutable in moe_forward registration. (#20152)
Signed-off-by: Bill Nell <bnell@redhat.com>
1 parent aafabaa commit c6c9830

File tree

1 file changed

+2
-1
lines changed
  • vllm/model_executor/layers/fused_moe

1 file changed

+2
-1
lines changed

vllm/model_executor/layers/fused_moe/layer.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1743,7 +1743,8 @@ def moe_forward_fake(hidden_states: torch.Tensor, router_logits: torch.Tensor,
17431743
direct_register_custom_op(
17441744
op_name="moe_forward",
17451745
op_func=moe_forward,
1746-
mutates_args=[],
1746+
mutates_args=["hidden_states"],
17471747
fake_impl=moe_forward_fake,
17481748
dispatch_key=current_platform.dispatch_key,
1749+
tags=(torch.Tag.needs_fixed_stride_order, ),
17491750
)

0 commit comments

Comments
 (0)