Update app.py

app.py

@@ -115,13 +115,24 @@ for i in range(n_start, n):
   mqa_time = shared_time + qkv_mqa_exec(bs, h, i, d)[2] + att1_mqa_exec(bs, h, i, d)[2] + att2_mqa_exec(bs, h, i, d)[2]
   mqa_total_time += l*mqa_time
   
-st.write("Multi-Head Attention: " + str(mha_total_time))
-st.write("Multi-Query Attention: " + str(mqa_total_time))
-st.write("Speed-up MQA over MHA: " + str(mha_total_time/mqa_total_time))
+c1, c2 = st.columns(2, 4)
+c1.write("Multi-Head Attention:")
+c2.write(str(round(mha_total_time, 2)))
+c1.write("Multi-Query Attention:")
+c2.write(str(round(mqa_total_time, 2)))
+c1.write("Speed-up MQA over MHA: ")
+c2.write(str(round(mha_total_time/mqa_total_time),2))
 
 st.header("Memory consumption")
+st.caption("MHA")
+c1, c2 = st.columns(2, 4)
 num_params = 12*l*d*d
-st.write("Num Parameters (in B)" + str(round(num_params/1e9, 3)))
+c1.write("Num Parameters (in B)")
+c2.write(str(round(num_params/1e9, 3)))
+c1.write("Storing activations")
+acts = round(2*l*(d/h)*h*n/1e9, 2)
+c2.write(str(acts))
+
 
 
 breakdown = st.checkbox("Show breakdown per operation")