Update app.py

app.py

@@ -95,7 +95,7 @@ def mlp_exec(bs, h, n, d):
   exec_time = calc_exec_time(flop, nbytes)
   return flop, nbytes, exec_time
   
-def print_kernel_execution(flop, mem_bytes):
+def print_kernel_execution(flop, nbytes):
   c1, c2 = st.columns([2, 3])
   exec_time = calc_exec_time(flop, nbytes, include_overhead=False)
   flop = round(flop/1e9, 2)
@@ -175,16 +175,12 @@ if breakdown:
   st.write("Showing calculation for the maximum sequence length (n)")
   
   st.caption("Multi-Head 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
-  c1, c2 = st.columns([2, 3])
-  att1_mha_time = print_kernel_execution(c1, c2, mha_flop, mha_bytes)
+  flop, nbytes, exec_time = att1_mha_exec(bs, h, n, d)
+  print_kernel_execution(flop, nbytes)
   
   st.caption("Multi-Query Attention")
-  mqa_flop = 2*bs*h*(d/h)*n
-  mqa_bytes = 2*bs*h*(d/h) + 2*bs*n*(d/h) + 2*bs*h*n
-  c1, c2 = st.columns([2, 3])
-  att1_mqa_time = print_kernel_execution(c1, c2, mqa_flop, mqa_bytes)
+  flop, nbytes, exec_time = att1_mqa_exec(bs, h, n, d)
+  print_kernel_execution(flop, nbytes)
   
   st.subheader('Attention-value gemm')
   st.write("Showing calculation for the maximum sequence length (n)")