add prefill memory

.streamlit/config.toml

@@ -0,0 +1 @@
+[theme]

app.py

@@ -57,13 +57,17 @@ subtotal_operations = [
 
 
 
-col1, col2, col3, col4, col5 = st.columns(5)
+col1, col2, col3, col4, col5 = st.columns([1,1.5,2,2,2])
 
 inference_config = {}
 parameter_count = {}
 cached_parameter_count = {}
+
 prefilling_operation_count = {}
 generation_operation_count = {}
+prefilling_memory_count = {}
+generation_memory_count = {}
+
 gpu_config = {}
 inference_info = {}
 
@@ -77,6 +81,7 @@ with col1:
     model_config['intermediate_size'] = st.number_input('intermediate size', value=model_config['intermediate_size'], format ="%d")
     model_config['vocab_size'] = st.number_input('vocab size', value= model_config['vocab_size'], format ="%d")
     model_config['max_position_embeddings'] = st.number_input('max position embeddings', value=model_config['max_position_embeddings'], format ="%d")
+    model_config['hidden_size_per_head'] = model_config['hidden_size']/model_config['num_attention_heads']
 
     header4("Inference Setting")
     inference_config['batchsize'] = st.number_input('batchsize', value=1, format ="%d")
@@ -131,43 +136,57 @@ with col2:
 
 with col3: # Prefilling
     prefilling_operation_count = prefilling_operation(model_config, inference_config)
-    inference_info['inference_prefilling_time'] = prefilling_operation_count['total'] / (gpu_config['TFLOP']*10**12)
+    prefilling_activation_memory_count = prefilling_activation_memory(model_config, inference_config)
+    inference_info['inference_prefilling_time'] = prefilling_operation_count['total'] / (gpu_config['TFLOP']*1024**4)
     inference_info['inference_prefilling_throughput'] = inference_config['input_seq_length']*inference_config['batchsize']/inference_info['inference_prefilling_time']
+    inference_info['prefilling_memory_latency'] = prefilling_activation_memory_count['total'] / (gpu_config['memory_bandwidth']*1024**3)
     cached_parameter_count['kv_cache'] = 2 * (inference_config['batchsize'] * (model_config['hidden_size'] * model_config['num_hidden_layers'] * inference_config['input_seq_length']))
 
     operation_items = {key: "{:,}".format(int(prefilling_operation_count[key])) for key in prefilling_operation_count if key not in subtotal_operations}
     subtotal_operation_items = {key: "{:,}".format(int(prefilling_operation_count[key])) for key in prefilling_operation_count if key in subtotal_operations}
+    prefilling_activation_memory_count = {key: "{:,}".format(int(value)) for key, value in prefilling_activation_memory_count.items()}
 
     ## Convert dictionaries to pandas dataframes for table display
     df_operation_count = pd.DataFrame(list(operation_items.items()), columns=["Operation", "FLOPS"])
     df_subtotal_operation_count = pd.DataFrame(list(subtotal_operation_items.items()), columns=["Operation", "FLOPS"])
-   
+    
+    df_operation_count["Activation (Byte)"] = df_operation_count["Operation"].map(prefilling_activation_memory_count)
+    df_subtotal_operation_count["Activation (Byte)"] = df_subtotal_operation_count["Operation"].map(prefilling_activation_memory_count)
+
     header4("Inference Ops: Prefilling")
     st.markdown(create_table(df_operation_count))
 
     header5("Summary: Prefilling")
     st.markdown(create_table(df_subtotal_operation_count))
     st.write(f"Prefillng throughput (tokens/s): {inference_info['inference_prefilling_throughput']:.2f}")
+    st.write(f"FLOPS latency: {inference_info['inference_prefilling_time']}")
+    st.write(f"Memory latency: {inference_info['prefilling_memory_latency']}")
 
     if inference_config['KV_cache']:
         st.write(f"kv cache (Byte): {cached_parameter_count['kv_cache']:,}")
         
 
 
-with col4: # Prefilling
+with col4: # Generation
     generation_operation_count = generation_operation(model_config, inference_config)
-    inference_info['inference_generation_time'] = generation_operation_count['total'] / (gpu_config['TFLOP']*10**12)
+    generation_activation_memory_count = generation_activation_memory(model_config, inference_config)
+    inference_info['inference_generation_time'] = generation_operation_count['total'] / (gpu_config['TFLOP']*1024**4)
     inference_info['inference_generation_throughput'] = inference_config['output_seq_length']*inference_config['batchsize']/inference_info['inference_generation_time']
     inference_info['inference_client_generation_throughput'] = inference_config['output_seq_length']*inference_config['batchsize'] / (inference_info['inference_prefilling_time'] + inference_info['inference_generation_time'])
+    inference_info['generation_memory_latency'] = generation_activation_memory_count['total'] / (gpu_config['memory_bandwidth']*1024**3)
     cached_parameter_count['kv_cache'] = 2 * (inference_config['batchsize'] * (model_config['hidden_size'] * model_config['num_hidden_layers'] * (inference_config['input_seq_length']+inference_config['output_seq_length'])))
 
     operation_items = {key: "{:,}".format(int(generation_operation_count[key])) for key in generation_operation_count if key not in subtotal_operations}
     subtotal_operation_items = {key: "{:,}".format(int(generation_operation_count[key])) for key in generation_operation_count if key in subtotal_operations}
+    generation_activation_memory_count = {key: "{:,}".format(int(value)) for key, value in generation_activation_memory_count.items()}
 
     ## Convert dictionaries to pandas dataframes for table display
     df_operation_count = pd.DataFrame(list(operation_items.items()), columns=["Operation", "FLOPS"])
     df_subtotal_operation_count = pd.DataFrame(list(subtotal_operation_items.items()), columns=["Operation", "FLOPS"])
    
+    #df_operation_count["Activation (Byte)"] = df_operation_count["Operation"].map(generation_activation_memory_count)
+    #df_subtotal_operation_count["Activation (Byte)"] = df_subtotal_operation_count["Operation"].map(generation_activation_memory_count)
+
     header4("Inference Ops: Generation")
     st.markdown(create_table(df_operation_count))
 
@@ -175,6 +194,8 @@ with col4: # Prefilling
     st.markdown(create_table(df_subtotal_operation_count))
     st.write(f"Generation-only throughput (tokens/s): {inference_info['inference_generation_throughput']:.2f}")
     st.write(f"(Client) Generation throughput (tokens/s): {inference_info['inference_client_generation_throughput']:.2f}")
+    st.write(f"FLOPS latency: {inference_info['inference_generation_time']}")
+    #st.write(f"Memory latency: {inference_info['generation_memory_latency']}")
 
     if inference_config['KV_cache']:
         st.write(f"kv cache (Byte): {cached_parameter_count['kv_cache']:,}")

calc_util.py

@@ -32,23 +32,23 @@ def positional_embedding_operation(model_config, inference_config):
 ### Below three are the same
 def attention_K_operation(model_config, inference_config, seq_length):
     A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
-    B = [model_config['hidden_size'], model_config['hidden_size']/model_config['num_attention_heads']]
+    B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
     return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * matrix_operation(A, B)
 
 def attention_Q_operation(model_config, inference_config, seq_length):
     A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
-    B = [model_config['hidden_size'], model_config['hidden_size']/model_config['num_attention_heads']]
+    B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
     return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * matrix_operation(A, B)
 
 def attention_V_operation(model_config, inference_config, seq_length):
     A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
-    B = [model_config['hidden_size'], model_config['hidden_size']/model_config['num_attention_heads']]
+    B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
     return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * matrix_operation(A, B)
 
 ## 
 def attention_QK_operation(model_config, inference_config, seq_length_Q, seq_length_K):
-    A = [inference_config['batchsize'], seq_length_Q, model_config['hidden_size']/model_config['num_attention_heads']]
-    B = [model_config['hidden_size']/model_config['num_attention_heads'], seq_length_K]
+    A = [inference_config['batchsize'], seq_length_Q, model_config['hidden_size_per_head']]
+    B = [model_config['hidden_size_per_head'], seq_length_K]
     return model_config['num_hidden_layers'] * model_config['num_attention_heads']* matrix_operation(A, B)
 
 def attention_softmax_operation(model_config, inference_config,seq_length):
@@ -59,7 +59,7 @@ def attention_softmax_operation(model_config, inference_config,seq_length):
 
 def attention_multV_operation(model_config, inference_config, seq_length_Q, seq_length_V):
     A = [inference_config['batchsize'], seq_length_Q, seq_length_V]
-    B = [seq_length_V, model_config['hidden_size']/model_config['num_attention_heads']]
+    B = [seq_length_V, model_config['hidden_size_per_head']]
     return model_config['num_hidden_layers'] * model_config['num_attention_heads']* matrix_operation(A, B)
 
 def attention_out_operation(model_config, inference_config, seq_length):
@@ -153,4 +153,133 @@ def generation_operation(model_config, inference_config):
     generation_operation_count['mlp'] = generation_operation_count['mlp1'] + generation_operation_count['mlp2']
     generation_operation_count['total'] = (generation_operation_count['attention'] + generation_operation_count['mlp'] + generation_operation_count['layernorm'])
 
-    return generation_operation_count
+    return generation_operation_count
+
+
+def word_embedding_activation_memory(model_config, inference_config, seq_length):
+    return inference_config['batchsize'] * seq_length * (model_config['vocab_size'] + model_config['hidden_size'])
+
+def positional_embedding_activation_memory(model_config, inference_config, seq_length):
+    return 2 * inference_config['batchsize'] * seq_length * model_config['hidden_size']
+
+def attention_K_activation_memory(model_config, inference_config, seq_length):
+    per_head_per_layer = inference_config['batchsize'] * seq_length * (model_config['hidden_size'] + model_config['hidden_size_per_head'])
+    return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * per_head_per_layer
+
+def attention_V_activation_memory(model_config, inference_config, seq_length):
+    per_head_per_layer = inference_config['batchsize'] * seq_length * (model_config['hidden_size'] + model_config['hidden_size_per_head'])
+    return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * per_head_per_layer
+
+def attention_Q_activation_memory(model_config, inference_config, seq_length):
+    per_head_per_layer = inference_config['batchsize'] * seq_length * (model_config['hidden_size'] + model_config['hidden_size_per_head'])
+    return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * per_head_per_layer
+
+def attention_QK_activation_memory(model_config, inference_config, seq_length_Q, seq_length_K):
+    inputs_Q = inference_config['batchsize'] * seq_length_Q * model_config['hidden_size_per_head']
+    inputs_K = inference_config['batchsize'] * seq_length_K * model_config['hidden_size_per_head']
+    outputs =  inference_config['batchsize'] * seq_length_Q * seq_length_K
+    return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * (inputs_Q + inputs_K + outputs)
+
+def attention_softmax_activation_memory(model_config, inference_config, seq_length):
+    per_head_per_layer = (2 * inference_config['batchsize'] * seq_length * seq_length)
+    return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * per_head_per_layer
+
+def attention_multV_activation_memory(model_config, inference_config, seq_length):
+    per_head_per_layer = inference_config['batchsize'] * seq_length * seq_length + 2 * inference_config['batchsize'] * seq_length * model_config['hidden_size_per_head']
+    return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * per_head_per_layer
+
+def attention_out_activation_memory(model_config, inference_config, seq_length):
+    per_head_per_layer = 2 * inference_config['batchsize'] * seq_length * model_config['hidden_size']
+    return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * per_head_per_layer
+
+def layernorm_activation_memory(model_config, inference_config, seq_length):
+    per_layernorm_per_layer = 2 * inference_config['batchsize'] * seq_length * model_config['hidden_size']
+    return model_config['num_hidden_layers'] * model_config['layernorm_operation'] * per_layernorm_per_layer
+
+def mlp1_activation_memory(model_config, inference_config, seq_length):
+    per_layer = inference_config['batchsize'] * seq_length * (model_config['hidden_size'] + model_config['intermediate_size'])
+    return model_config['num_hidden_layers'] * per_layer
+
+def mlp2_activation_memory(model_config, inference_config, seq_length):
+    per_layer = inference_config['batchsize'] * seq_length * (model_config['intermediate_size'] + model_config['hidden_size'])
+    return model_config['num_hidden_layers'] * per_layer
+
+def prefilling_activation_memory(model_config, inference_config):
+    activation_memory = {}
+    
+    activation_memory['word_embedding'] = word_embedding_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+    activation_memory['positional_embedding'] = positional_embedding_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+    
+    activation_memory['attention_Q'] = attention_Q_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+    activation_memory['attention_K'] = attention_K_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+    activation_memory['attention_V'] = attention_V_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+    activation_memory['attention_QK'] = attention_QK_activation_memory(model_config, inference_config, inference_config['input_seq_length'], inference_config['input_seq_length'])
+    activation_memory['attention_softmax'] = attention_softmax_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+    activation_memory['attention_multV'] = attention_multV_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+    activation_memory['attention_out'] = attention_out_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+
+    activation_memory['layernorm'] = layernorm_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+
+    activation_memory['mlp1'] = mlp1_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+    activation_memory['mlp2'] = mlp2_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+    
+    activation_memory['embeddings'] = activation_memory['word_embedding'] + activation_memory['positional_embedding']
+    activation_memory['attention'] = (
+        activation_memory['attention_Q'] + activation_memory['attention_K'] +
+        activation_memory['attention_V'] + activation_memory['attention_QK'] +
+        activation_memory['attention_softmax'] + activation_memory['attention_multV'] +
+        activation_memory['attention_out']
+    )
+    activation_memory['mlp'] = activation_memory['mlp1'] + activation_memory['mlp2']
+    activation_memory['total'] = (
+        activation_memory['embeddings'] + activation_memory['attention'] +
+        activation_memory['mlp'] + activation_memory['layernorm']
+    )
+    
+    activation_memory['embeddings'] = activation_memory['word_embedding'] + activation_memory['positional_embedding'] 
+    activation_memory['attention'] = sum([v for k,v in activation_memory.items() if 'attention' in k])
+    activation_memory['mlp'] = activation_memory['mlp1'] + activation_memory['mlp2']
+    activation_memory['total'] = (activation_memory['attention'] + activation_memory['mlp'] + activation_memory['layernorm'])
+
+    return activation_memory
+
+
+def generation_activation_memory(model_config, inference_config):
+    # TODO Check how KV cache affects activation_memory
+    activation_memory = {}
+    
+    activation_memory['word_embedding'] = word_embedding_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'])
+    activation_memory['positional_embedding'] = positional_embedding_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'])
+    
+    activation_memory['attention_Q'] = attention_Q_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'])
+    activation_memory['attention_K'] = attention_K_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'])
+    activation_memory['attention_V'] = attention_V_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'])
+    activation_memory['attention_QK'] = attention_QK_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'], inference_config['input_seq_length'] + inference_config['output_seq_length'])
+    activation_memory['attention_softmax'] = attention_softmax_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'])
+    activation_memory['attention_multV'] = attention_multV_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'])
+    activation_memory['attention_out'] = attention_out_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'])
+
+    activation_memory['layernorm'] = layernorm_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'])
+
+    activation_memory['mlp1'] = mlp1_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'])
+    activation_memory['mlp2'] = mlp2_activation_memory(model_config, inference_config, inference_config['input_seq_length'] + inference_config['output_seq_length'])
+    
+    activation_memory['embeddings'] = activation_memory['word_embedding'] + activation_memory['positional_embedding']
+    activation_memory['attention'] = (
+        activation_memory['attention_Q'] + activation_memory['attention_K'] +
+        activation_memory['attention_V'] + activation_memory['attention_QK'] +
+        activation_memory['attention_softmax'] + activation_memory['attention_multV'] +
+        activation_memory['attention_out']
+    )
+    activation_memory['mlp'] = activation_memory['mlp1'] + activation_memory['mlp2']
+    activation_memory['total'] = (
+        activation_memory['embeddings'] + activation_memory['attention'] +
+        activation_memory['mlp'] + activation_memory['layernorm']
+    )
+    
+    activation_memory['embeddings'] = activation_memory['word_embedding'] + activation_memory['positional_embedding'] 
+    activation_memory['attention'] = sum([v for k,v in activation_memory.items() if 'attention' in k])
+    activation_memory['mlp'] = activation_memory['mlp1'] + activation_memory['mlp2']
+    activation_memory['total'] = (activation_memory['attention'] + activation_memory['mlp'] + activation_memory['layernorm'])
+
+    return activation_memory