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

sgl-project · Nov 24, 2024 · fa27161 · fa27161
1 parent 5652c56
commit fa27161
Showing 1 changed file with 3 additions and 2 deletions.
diff --git 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