Update app.py

app.py

@@ -28,6 +28,9 @@ else:
 # in ms
 
 THREAD_OVERHEAD = st.sidebar.number_input('Thread overhead (in ms)', format="%.3f", value=0.005)
+GPU_EFFICIENCY = st.sidebar.number_input('GPU efficiency', format="%.3f", value=0.5)
+
+TFLOPS = GPU_EFFICIENCY*TFLOPS
 
 # in ms
 def calc_exec_time(comp_flop, mem_bytes, include_overhead=True):
@@ -112,7 +115,8 @@ def print_kernel_execution(flop, nbytes):
   c1.write("Overhead (ms):")
   c2.write(str(THREAD_OVERHEAD))
 
-st.header("Execution time (ms)")
+st.title("Inference time MHA vs MQA")
+st.write("This space approximates the inference time for Multi-Query Attention and Multi-Head Attention model. You can change the hyperparameters in sidebar.")
 
 mqa_total_time = 0.
 mha_total_time = 0.
@@ -134,7 +138,7 @@ c1.write("Speed-up MQA over MHA: ")
 c2.write(str(round(mha_total_time/mqa_total_time,2)))
 
 st.header("Memory consumption")
-st.caption("MHA")
+st.subheader("MHA")
 c1, c2 = st.columns([2, 4])
 num_params = 12*l*d*d
 c1.write("Num Parameters (in B)")
@@ -145,7 +149,7 @@ c1.write("Cached keys and values (GB)")
 acts = round(2*bs*l*(d/h)*h*2*n/1e9, 2)
 c2.write(str(acts))
 
-st.caption("MQA")
+st.subheader("MQA")
 c1, c2 = st.columns([2, 4])
 num_params = (10+2/h)*l*d*d
 c1.write("Num Parameters (in B)")
@@ -156,6 +160,8 @@ c1.write("Cached keys and values (GB)")
 acts = round(2*bs*l*(d/h)*2*n/1e9, 2)
 c2.write(str(acts))
 
+st.subheader("How we approximate the inference time")
+st.write("TODO")
 
 breakdown = st.checkbox("Show breakdown per operation")
 if breakdown: