Update app.py

app.py

@@ -137,8 +137,8 @@ c2.write(str(round(mqa_total_time, 2)))
 c1.write("Speed-up MQA over MHA: ")
 c2.write(str(round(mha_total_time/mqa_total_time,2)))
 
-st.header("Memory consumption")
-st.subheader("MHA")
+st.subheader("Memory consumption")
+st.caption("Multi-Head Attention")
 c1, c2 = st.columns([2, 4])
 num_params = 12*l*d*d
 c1.write("Num Parameters (in B)")
@@ -149,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.subheader("MQA")
+st.subheader("Multi-Query Attention")
 c1, c2 = st.columns([2, 4])
 num_params = (10+2/h)*l*d*d
 c1.write("Num Parameters (in B)")