app.py

import streamlit as st

st.header("Transformer parameters")
bs = st.number_input('Batch size: ', value=10)
h = st.number_input('Num heads: ', value=16)
d = st.number_input('Dimension: ', value=768)
n = st.number_input('Seq length: ', value=1024)

st.header('Query, Key, Value projection')

mha_flop = 2*bs*n*d*3*d
mha_bytes = 2*bs*n*d + 2*3*d*d + 2*bs*n*3*d

st.subheader("Multi-query Attention")
st.text("FLOP: " + str(mha_flop))
st.text("bytes: " + str(mha_bytes))
st.text("Arithm. intensity:" + str(mha_flop/mha_bytes))

mqa_flop = 2*bs*n*d*(1+2/h)*d
mqa_bytes = 2*bs*n*d + 2*(2/h)*d*d + 2*bs*n*(2/h)*d

st.subheader("Multi-query Attention")
st.write("FLOP: " + str(mqa_flop))
st.write("bytes: " + str(mqa_bytes))
st.write("Arithm. intensity:" + str(mqa_flop/mqa_bytes))

st.header('Attention')