add generation arithmetic intensity

app.py

@@ -57,7 +57,7 @@ subtotal_operations = [
 
 
 
-col1, col2, col3, col4, col5 = st.columns([1,1.5,2.3,2.3,0.1])
+col1, col2, col3, col4, col5 = st.columns([1,1.5,2.5,2.5,0.1])
 
 inference_config = {}
 parameter_count = {}
@@ -144,7 +144,7 @@ with col3: # Prefilling
 
     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_arithmetic_intensity = {key: "{:.3f}".format(prefilling_operation_count[key]/prefilling_activation_memory_count[key]) for key in prefilling_activation_memory_count}
+    prefilling_arithmetic_intensity = {key: "{:.3f}".format(prefilling_operation_count[key]/prefilling_activation_memory_count[key] if prefilling_activation_memory_count[key]>0 else float('inf'))  for key in prefilling_activation_memory_count}
     prefilling_activation_memory_count = {key: "{:,}".format(int(value)) for key, value in prefilling_activation_memory_count.items()}
 
 
@@ -182,15 +182,18 @@ with col4: # Generation
 
     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_arithmetic_intensity = {key: "{:.3f}".format(generation_operation_count[key]/generation_activation_memory_count[key] if generation_activation_memory_count[key]>0 else float('inf')) for key in generation_activation_memory_count}
     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)
-
+    df_operation_count["Activation (Byte)"] = df_operation_count["Operation"].map(generation_activation_memory_count)
+    df_operation_count["Arithmetic Intensity"] = df_operation_count["Operation"].map(generation_arithmetic_intensity)
+    df_subtotal_operation_count["Activation (Byte)"] = df_subtotal_operation_count["Operation"].map(generation_activation_memory_count)
+    df_subtotal_operation_count["Arithmetic Intensity"] = df_subtotal_operation_count["Operation"].map(generation_arithmetic_intensity)
+    
     header4("Inference Ops: Generation")
     st.markdown(create_table(df_operation_count))
 
@@ -199,7 +202,7 @@ with col4: # Generation
     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']}")
+    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

@@ -184,8 +184,8 @@ def attention_softmax_activation_memory(model_config, inference_config, seq_leng
     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']
+def attention_multV_activation_memory(model_config, inference_config, seq_length_Q, seq_length_V):
+    per_head_per_layer = inference_config['batchsize'] * seq_length_Q * seq_length_V + inference_config['batchsize'] * seq_length_Q * model_config['hidden_size_per_head'] + inference_config['batchsize'] * seq_length_V * 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):
@@ -215,7 +215,7 @@ def prefilling_activation_memory(model_config, inference_config):
     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_multV'] = attention_multV_activation_memory(model_config, inference_config, inference_config['input_seq_length'], 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'])
@@ -243,31 +243,50 @@ def prefilling_activation_memory(model_config, inference_config):
 
     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['word_embedding'] = 0
+    activation_memory['positional_embedding'] = 0
+    activation_memory['attention_K'] = 0
+    activation_memory['attention_V'] = 0
+    activation_memory['attention_Q'] = 0
+    activation_memory['attention_QK'] = 0
+    activation_memory['attention_softmax'] = 0
+    activation_memory['attention_multV'] = 0
+    activation_memory['attention_out'] = 0
+    activation_memory['mlp1'] = 0
+    activation_memory['mlp2'] = 0
+    activation_memory['layernorm'] = 0
+
+    for t in range(inference_config['output_seq_length']):
+        if inference_config['KV_cache']:
+            activation_memory['attention_K'] += attention_K_activation_memory(model_config, inference_config, 1)
+            activation_memory['attention_V'] += attention_V_activation_memory(model_config, inference_config, 1)
+            activation_memory['attention_Q'] += attention_Q_activation_memory(model_config, inference_config, 1)
+            activation_memory['attention_QK'] += attention_QK_activation_memory(model_config, inference_config, seq_length_Q=1, seq_length_K=(t+1)+inference_config['input_seq_length'])
+            activation_memory['attention_softmax'] += attention_softmax_activation_memory(model_config, inference_config, 1)
+            activation_memory['attention_multV'] += attention_multV_activation_memory(model_config, inference_config, seq_length_Q=1, seq_length_V=(t+1)+inference_config['input_seq_length'])
+            activation_memory['attention_out'] += attention_out_activation_memory(model_config, inference_config, 1)
+            activation_memory['mlp1'] += mlp1_activation_memory(model_config, inference_config, 1)
+            activation_memory['mlp2'] += mlp2_activation_memory(model_config, inference_config, 1)
+        else:
+            activation_memory['attention_K'] += attention_K_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
+            activation_memory['attention_V'] += attention_V_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
+            activation_memory['attention_Q'] += attention_Q_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
+            activation_memory['attention_QK'] += attention_QK_activation_memory(model_config, inference_config, seq_length_Q=(t+1)+inference_config['input_seq_length'], seq_length_K=(t+1)+inference_config['input_seq_length'])
+            activation_memory['attention_softmax'] += attention_softmax_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
+            activation_memory['attention_multV'] += attention_multV_activation_memory(model_config, inference_config, seq_length_Q=(t+1)+inference_config['input_seq_length'], seq_length_V=(t+1)+inference_config['input_seq_length'])
+            activation_memory['attention_out'] += attention_out_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
+            activation_memory['mlp1'] += mlp1_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
+            activation_memory['mlp2'] += mlp2_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
+
+        activation_memory['layernorm'] += layernorm_activation_memory(model_config, inference_config, (t+1)+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_K'] + activation_memory['attention_V'] +
+        activation_memory['attention_Q'] + activation_memory['attention_QK'] +
         activation_memory['attention_softmax'] + activation_memory['attention_multV'] +
         activation_memory['attention_out']
     )
@@ -276,10 +295,5 @@ def generation_activation_memory(model_config, inference_config):
         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