Update app.py

app.py

@@ -28,7 +28,7 @@ mha_flop = round(mha_flop/1e9, 2)
 mha_bytes = round(mha_bytes/1e6, 2)
 
 
-st.subheader("Multi-query Attention")
+st.subheader("Multi-head Attention")
 c1, c2 = st.columns([2, 3])
 c1.write("GFLOP:")
 c2.write(str(mha_flop))
@@ -42,6 +42,8 @@ c2.write(str(mha_time))
 
 mqa_flop = 2*bs*n*d*(1+2/h)*d
 mqa_bytes = 2*bs*n*d + 2*(2/h)*d*d + 2*bs*n*(2/h)*d
+mqa_intensity = mqa_flop/mqa_bytes
+mqa_time = (mqa_flop/TFLOPS + mqa_bytes/GB_S)*1000
 
 mqa_flop = round(mqa_flop/1e9, 2)
 mqa_bytes = round(mqa_bytes/1e6, 2)
@@ -53,7 +55,9 @@ c2.write(str(mqa_flop))
 c1.write("MB:")
 c2.write(str(mqa_bytes))
 c1.write("Arithm. intensity:")
-c2.write(str(mqa_flop/mqa_bytes))
+c2.write(str(mqa_intensity))
+c1.write("Time (ms):")
+c2.write(str(mqa_time))
 
 st.header('Attention')