Update app.py

app.py

@@ -19,8 +19,9 @@ n = number_field('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
+
+mha_flop = 2*bs*1*d*3*d
+mha_bytes = 2*bs*1*d + 2*3*d*d + 2*bs*1*3*d
 mha_int = mha_flop/mha_bytes
 mha_time = (mha_flop/TFLOPS + mha_bytes/GB_S)*1000
 
@@ -28,7 +29,7 @@ mha_flop = round(mha_flop/1e9, 2)
 mha_bytes = round(mha_bytes/1e6, 2)
 
 
-st.subheader("Multi-head Attention")
+st.subheader("Multi-Head Attention")
 c1, c2 = st.columns([2, 3])
 c1.write("GFLOP:")
 c2.write(str(mha_flop))
@@ -40,15 +41,15 @@ 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 = 2*bs*1*d*(1+2/h)*d
+mqa_bytes = 2*bs*1*d + 2*(2/h)*d*d + 2*bs*1*(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)
 
-st.subheader("Multi-query Attention")
+st.subheader("Multi-Query Attention")
 c1, c2 = st.columns([2, 3])
 c1.write("GFLOP:")
 c2.write(str(mqa_flop))
@@ -60,4 +61,41 @@ c1.write("Time (ms):")
 c2.write(str(mqa_time))
 
 st.header('Attention')
+mha_flop = 2*bs*h*(d/h)*n
+mha_bytes = 2*bs*h*(d/h) + 2*bs*h*n*(d/h) + 2*bs*h*n*n
+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)
+
+
+st.subheader("Multi-Head Attention")
+c1, c2 = st.columns([2, 3])
+c1.write("GFLOP:")
+c2.write(str(mha_flop))
+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*h*(d/h)*n
+mqa_bytes = 2*bs*h*(d/h) + 2*bs*n*(d/h) + 2*bs*h*n*n
+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)
+
+st.subheader("Multi-Query Attention")
+c1, c2 = st.columns([2, 3])
+c1.write("GFLOP:")
+c2.write(str(mqa_flop))
+c1.write("MB:")
+c2.write(str(mqa_bytes))
+c1.write("Arithm. intensity:")
+c2.write(str(mqa_intensity))
+c1.write("Time (ms):")
+c2.write(str(mqa_time))