Update app.py

app.py

@@ -6,6 +6,9 @@ def number_field(label, **args):
   
   return c2.number_input('', **args)
 
+TFLOPS = 312e12
+GB_S = 1935e9
+
 st.header("Transformer parameters")
 col1, col2 = st.columns([2, 4])
 
@@ -19,6 +22,7 @@ st.header('Query, Key, Value projection')
 mha_flop = 2*bs*n*d*3*d
 mha_bytes = 2*bs*n*d + 2*3*d*d + 2*bs*n*3*d
 mha_int = mha_flop/mha_bytes
+mha_time = (mha_flop/TFLOPS + mha_bytes/GB_S)*1000
 
 mha_flop = round(mha_flop/1e9, 2)
 mha_bytes = round(mha_bytes/1e6, 2)
@@ -32,18 +36,24 @@ c1.write("MB: ")
 c2.write(str(mha_bytes))
 c1.write("Arithm. intensity:")
 c2.write(str(mha_int))
+c1.write("Time (ms):")
+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_flop = round(mqa_flop/1e9, 2)
+mqa_bytes = round(mqa_bytes/1e6, 2)
+
 st.subheader("Multi-query Attention")
 c1, c2 = st.columns([2, 3])
-c1.write("FLOP:")
+c1.write("GFLOP:")
 c2.write(str(mqa_flop))
-c1.write("Bytes: ")
+c1.write("MB:")
 c2.write(str(mqa_bytes))
 c1.write("Arithm. intensity:")
 c2.write(str(mqa_flop/mqa_bytes))
 
 st.header('Attention')
+