[optimization] remove python function call for custom op (#11750)

Signed-off-by: youkaichao <youkaichao@gmail.com>
2025-01-08 01:04:28 +08:00
parent c0efe92d8b
commit 869579a702
4 changed files with 15 additions and 13 deletions
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -753,7 +753,8 @@ def fused_experts_impl(hidden_states: torch.Tensor,
                                use_int8_w8a16=use_int8_w8a16,
                                block_shape=block_shape)

-        ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        torch.ops._C.silu_and_mul(intermediate_cache2,
+                                  intermediate_cache1.view(-1, N))

        invoke_fused_moe_kernel(intermediate_cache2,
                                w2,