Update app.py

app.py

@@ -53,7 +53,7 @@ st.caption("Multi-Query Attention")
 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
 c1, c2 = st.columns([2, 3])
-qkv_mha_time = print_kernel_execution(c1, c2, mqa_flop, mqa_bytes)
+qkv_mqa_time = print_kernel_execution(c1, c2, mqa_flop, mqa_bytes)
 
 st.subheader('QK gemm')
 st.write("Note that calculation depends on sequence length (n)")
@@ -101,7 +101,7 @@ softmax_time = print_kernel_execution(c1, c2, 0, softmax_bytes)
 st.caption("Layer norm/residual connection")
 ln_bytes = 2*bs*1*d
 ln_flop = 0
-ln_time = print_kernel_execution(c1, c2, 0, softmax_bytes)
+ln_time = print_kernel_execution(c1, c2, 0, ln_bytes)
 
 st.header('MLP')
 st.subheader('First Linear')
@@ -120,7 +120,7 @@ st.subheader('Element-wise ops')
 st.write("We also need to take into the GeLU, layer norm, and residual connection. We assume that these operations are memory bound. ")
 ln_bytes = 2*bs*1*d
 ln_flop = 0
-ln_time = print_kernel_execution(c1, c2, 0, softmax_bytes)
+ln_time = print_kernel_execution(c1, c2, 0, ln_bytes)
 
 st.header("Adding it all up")