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

fa271613 · Yineng Zhang · GitHub · 5652c565 · fa271613
Unverified Commit fa271613 authored Nov 24, 2024 by Yineng Zhang Committed by GitHub Nov 24, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 2 deletions

python/sglang/srt/layers/triton_fused_moe/fused_moe.py python/sglang/srt/layers/triton_fused_moe/fused_moe.py +3 -2

No files found.
--- 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