Update app.py

app.py

@@ -34,9 +34,13 @@ mha_flop = 2*bs*n*d*3*d
 mha_bytes = 2*bs*n*d + 2*3*d*d + 2*bs*n*3*d
 
 st.subheader("Multi-query Attention")
-key_value("FLOP: ", str(mha_flop))
-key_value("bytes: ", str(mha_bytes))
-key_value("Arithm. intensity:", str(mha_flop/mha_bytes))
+c1, c2 = st.columns(columns or [2, 3])
+c1.write("FLOP:")
+c2.write(str(mha_flop))
+c1.write("Bytes: ")
+c2.write(str(mha_bytes))
+c1.write("Arithm. intensity:")
+c2.write(str(mha_flop/mha_bytes))
 
 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