import streamlit as st

st.sidebar.header("Transformer parameters")
col1, col2 = st.sidebar.columns([2, 4])

bs = st.sidebar.number_input('Batch size', value=10)
h = st.sidebar.number_input('Num heads',value=16)
d = st.sidebar.number_input('Dimension', value=768)
l = st.sidebar.number_input('Num layers', value=24)

n_start = st.sidebar.number_input('Start seq', value=1)
n = st.sidebar.number_input('End seq', value=1024)

st.sidebar.header("GPU parameters")

GPU = st.sidebar.selectbox('GPU', ('A100', 'V100'))

if GPU == 'A100':
  # A100 specs
  TFLOPS = 312e12 
  GB_S = 1935e9
elif GPU == 'V100':
  TFLOPS = 112e12
  GB_S = 900e9
else:
  raise ValueError('Unknown GPU') 

# 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):
  exec_time = max(comp_flop/TFLOPS, mem_bytes/GB_S)
  exec_time *= 1000
  if include_overhead:
    exec_time = max(exec_time, THREAD_OVERHEAD)
  return exec_time

def qkv_mha_exec(bs, h, n, d):
  flop = 2*bs*1*d*3*d
  nbytes = 2*bs*1*d + 2*3*d*d + 2*bs*1*3*d
  exec_time = calc_exec_time(flop, nbytes)
  return flop, nbytes, exec_time
     
def qkv_mqa_exec(bs, h, n, d):
  flop = 2*bs*1*d*(1+2/h)*d
  nbytes = 2*bs*1*d + 2*(2/h)*d*d + 2*bs*1*(2/h)*d
  exec_time = calc_exec_time(flop, nbytes)
  return flop, nbytes, exec_time
  
def att1_mha_exec(bs, h, n, d):
  flop = 2*bs*h*(d/h)*n
  nbytes = 2*bs*h*(d/h) + 2*bs*h*n*(d/h) + 2*bs*h*n
  exec_time = calc_exec_time(flop, nbytes)
  return flop, nbytes, exec_time
  
def att1_mqa_exec(bs, h, n, d):
  flop = 2*bs*h*(d/h)*n
  nbytes = 2*bs*h*(d/h) + 2*bs*n*(d/h) + 2*bs*h*n
  exec_time = calc_exec_time(flop, nbytes)
  return flop, nbytes, exec_time

def att2_mha_exec(bs, h, n, d):
  flop = 2*bs*h*n*(d/h)
  nbytes = 2*bs*h*n + 2*bs*h*n*(d/h) + 2*bs*h*(d/h)
  exec_time = calc_exec_time(flop, nbytes)
  return flop, nbytes, exec_time
  
def att2_mqa_exec(bs, h, n, d):
  flop = 2*bs*h*n*(d/h)
  nbytes = 2*bs*n*(d/h) + 2*bs*n*(d/h) + 2*bs*h*(d/h)
  exec_time = calc_exec_time(flop, nbytes)
  return flop, nbytes, exec_time
  
def out_exec(bs, h, n, d):
  flop = 2*bs*1*d*d
  nbytes = 2*bs*1*d + 2*d*d + 2*bs*1*d
  exec_time = calc_exec_time(flop, nbytes)
  return flop, nbytes, exec_time

def softmax_exec(bs, h, n, d):
  flop = 0
  nbytes = 2*bs*h*n + 2*bs*h*n
  exec_time = calc_exec_time(flop, nbytes)
  return flop, nbytes, exec_time

def ln_exec(bs, h, n, d):
  nbytes = 2*bs*1*d + 2*bs*1*d
  flop = 0
  exec_time = calc_exec_time(flop, nbytes)
  return flop, nbytes, exec_time

def mlp_exec(bs, h, n, d):
  flop = 2*bs*1*d*4*d
  nbytes = 2*bs*1*d + 2*d*4*d + 2*bs*1*4*d
  exec_time = calc_exec_time(flop, nbytes)
  return flop, nbytes, exec_time
  
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)
  nbytes = round(nbytes/1e6, 2)
  
  c1.write("GFLOP:")
  c2.write(str(flop))
  c1.write("MB: ")
  c2.write(str(nbytes))
  c1.write("Time (ms):")
  c2.write(str(exec_time))
  c1.write("Overhead (ms):")
  c2.write(str(THREAD_OVERHEAD))

st.title("Inference time MHA vs MQA")
st.write("This space approximates the inference time for Multi-Query Attention and Multi-Head Attention transformers. You can change the hyperparameters in sidebar.")

mqa_total_time = 0.
mha_total_time = 0.

for i in range(n_start, n):
  shared_time = out_exec(bs, h, i, d)[2] + softmax_exec(bs, h, i , d)[2] + 2*ln_exec(bs, h, i, d)[2] \
                + 2*mlp_exec(bs, h, i, d)[2] + 3*ln_exec(bs, h, i, d)[2]
  mha_time = shared_time + qkv_mha_exec(bs, h, i, d)[2] + att1_mha_exec(bs, h, i, d)[2] + att2_mha_exec(bs, h, i, d)[2]
  mha_total_time += l*mha_time
  mqa_time = shared_time + qkv_mqa_exec(bs, h, i, d)[2] + att1_mqa_exec(bs, h, i, d)[2] + att2_mqa_exec(bs, h, i, d)[2]
  mqa_total_time += l*mqa_time
  
c1, c2 = st.columns([2, 4])
c1.write("Multi-Head Attention:")
c2.write(str(round(mha_total_time, 2)))
c1.write("Multi-Query Attention:")
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.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)")
c2.write(str(round(num_params/1e9, 3)))
c1.write("Stored Parameters (GB)")
c2.write(str(round(2*num_params/1e9, 3)))
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("Multi-Query Attention")
c1, c2 = st.columns([2, 4])
num_params = (10+2/h)*l*d*d
c1.write("Num Parameters (in B)")
c2.write(str(round(num_params/1e9, 3)))
c1.write("Stored Parameters (GB)")
c2.write(str(round(2*num_params/1e9, 3)))
c1.write("Cached keys and values (GB)")
acts = round(2*bs*l*(d/h)*2*n/1e9, 2)
c2.write(str(acts))

st.subheader("Estimating execution time")
st.markdown("We use the [following crude approximation](https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html#understand-perf) to estimate the execution time for each matrix multiplication.")

st.latex("C = A \cdot B")
st.latex("A \in \mathbb{R}^{MxK}, B \in R^{KxN}, C \in \mathbb{R}^{MxN}")

st.markdown('''
To execute this operation on the GPU, we need to 
1. Read A, B from memory
2. Perform matrix multiplication
3. Write C to memory
''')

st.markdown("For float16 operations (2 bytes), we can estimate the memory access time of A as follows:")
st.latex("T_{mem}(A) = 2*M*K / BW_{mem}")
st.markdown("where BW_mem is the memory bandwidth of the GPU (e.g. 1935 GB/s for an A100 GPU)")
st.markdown("The total time on memory access is T_mem = T_mem(A) + T_mem(B) + T_mem(C)")

st.markdown("We can estimate the compute time for the math operations as follows:")
st.latex("T_{math}(A \cdot B) = 2*M*K*N / BW_{math}")
st.markdown("where BW_math is the number of floating point operations per second (e.g. 312 TFLOPS for an A100 GPU)")

st.markdown("If we assume we can *perfectly* overlap memory access with math operations, then the estimated execution time for the operation is:")
st.latex("max(T_{math}, T_{mem})")

st.markdown("Note that there is a minimum time to execute the operation due to [kernel launch overhead](https://forums.developer.nvidia.com/t/any-way-to-measure-the-latency-of-a-kernel-launch/221413/2)")

st.subheader("Inference time for Transformer operations")
st.markdown("We can now estimate the execution for each of the operations in the transformer model. I suggest you inspect the code for details on the calculations. ")

st.subheader('Attention layer')

st.markdown('**QKV projection**')
st.caption("Multi-Head Attention")
flop, nbytes, exec_time = qkv_mha_exec(bs, h, n, d)
print_kernel_execution(flop, nbytes)
 
st.caption("Multi-Query Attention")
flop, nbytes, exec_time = qkv_mqa_exec(bs, h, n, d)
print_kernel_execution(flop, nbytes)

st.markdown('**QK gemm**')
st.write("Showing calculation for the maximum sequence length (n)")

st.caption("Multi-Head Attention")
flop, nbytes, exec_time = att1_mha_exec(bs, h, n, d)
print_kernel_execution(flop, nbytes)

st.caption("Multi-Query Attention")
flop, nbytes, exec_time = att1_mqa_exec(bs, h, n, d)
print_kernel_execution(flop, nbytes)

st.markdown('**Attention-value gemm**')
st.write("Showing calculation for the maximum sequence length (n)")
st.caption("Multi-Head Attention")
flop, nbytes, exec_time = att2_mha_exec(bs, h, n, d)
print_kernel_execution(flop, nbytes)

st.caption("Multi-Query Attention")
flop, nbytes, exec_time = att2_mqa_exec(bs, h, n, d)
print_kernel_execution(flop, nbytes)

st.markdown('**Output projection**')
flop, nbytes, exec_time = out_exec(bs, h, n, d)
print_kernel_execution(flop, nbytes)

st.markdown('**Element-wise ops**')
st.write("We also need to take into the softmax layer, layer norm, and residual connection. We assume that these operations are memory bound. ")

st.caption("Softmax")
flop, nbytes, exec_time = softmax_exec(bs, h, n, d)
print_kernel_execution(flop, nbytes)

st.caption("Layer norm/residual connection")
flop, nbytes, exec_time = ln_exec(bs, h, n, d)
print_kernel_execution(flop, nbytes)

st.subheader('MLP layer')  
st.markdown('**First and Second Linear Layer**')
flop, nbytes, exec_time = mlp_exec(bs, h, n, d)
print_kernel_execution(flop, nbytes)

st.markdown('**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. ")
flop, nbytes, exec_time = ln_exec(bs, h, n, d)
print_kernel_execution(flop, nbytes)