Update app.py

app.py

@@ -6,9 +6,12 @@ def number_field(label, **kwargs):
   
   return c2.number_input('', **kwargs)
 
+def calc_exec_time(comp_flop, mem_bytes):
+  return (comp_flop/TFLOPS + mem_bytes/GB_S)*1000
+  
 def print_kernel_execution(c1, c2, comp_flop, mem_bytes):
   arith_int = comp_flop/mem_bytes
-  exec_time = (comp_flop/TFLOPS + mem_bytes/GB_S)*1000
+  exec_time = calc_exec_time(comp_flop, mem_bytes)
 
   comp_flop = round(comp_flop/1e9, 2)
   mem_bytes = round(mem_bytes/1e6, 2)
@@ -82,13 +85,21 @@ c1, c2 = st.columns([2, 3])
 att_mqa_time = print_kernel_execution(c1, c2, mqa_flop, mqa_bytes)
 
 st.subheader('Output projection')
-mlp1_flop = 2*bs*1*d
-mlp1_bytes = 2*bs*1*d + 2*d*4*d + 2*bs*1*4*d
+out_flop = 2*bs*1*d*d
+out_bytes = 2*bs*1*d + 2*d*d + 2*bs*1*d
 c1, c2 = st.columns([2, 3])
-mlp1_time = print_kernel_execution(c1, c2, mlp1_flop, mlp1_bytes)
+out_time = print_kernel_execution(c1, c2, out_flop, out_bytes)
 
 st.subheader('Element-wise ops')
-st.write("A couple of layers ")
+st.write("We also need to take into the softmax layer and layer norm")
+
+st.caption("Softmax")
+softmax_bytes = 2*bs*h*n + 2*bs*h*n
+c1, c2 = st.columns([2, 3])
+softmax_time = print_kernel_execution(c1, c2, 0, softmax_bytes))
+
+st.caption("Layer norm")
+
 
 st.header('MLP')
 st.subheader('First Linear')