fix: use torch.sum for compatible (#2161)

2024-11-24 22:37:04 +08:00
parent 5652c56535
commit fa27161380
1 changed files with 3 additions and 2 deletions
--- a/python/sglang/srt/layers/triton_fused_moe/fused_moe.py
+++ b/python/sglang/srt/layers/triton_fused_moe/fused_moe.py
@@ -766,9 +766,10 @@ def fused_experts_impl(
            use_int8_w8a16=use_int8_w8a16,
        )

-        ops.moe_sum(
+        torch.sum(
            intermediate_cache3.view(*intermediate_cache3.shape),
-            out_hidden_states[begin_chunk_idx:end_chunk_idx],
+            dim=1,
+            out=out_hidden_states[begin_chunk_idx:end_chunk_idx],
        )
    return out_hidden_states