Update app.py

app.py

@@ -7,6 +7,21 @@ d = st.number_input('Dimension: ', value=768)
 n = st.number_input('Seq length: ', value=1024)
 
 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
-st.write()
+
+st.write("Multi-query Attention")
+st.write("FLOP: " + mha_flop)
+st.write("bytes: " + mha_bytes)
+st.write("Arithm. intensity:" + (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
+
+st.write("Multi-query Attention")
+st.write("FLOP: " + mqa_flop)
+st.write("bytes: " + mqa_bytes)
+st.write("Arithm. intensity:" + (mqa_flop/mqa_bytes))
+
+st.header('Query, Key, Value projection')