Simplify ncu scripts.

microsoft · Sep 19, 2024 · 2264119 · 2264119
1 parent fc061ce
commit 2264119
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 8 deletions.
diff --git a/artifacts/run_all_ncu_cutlass.sh b/artifacts/run_all_ncu_cutlass.sh
@@ -17,10 +17,11 @@ fi
 # 1. ncu test the mha benchmark
 echo "NCU profiling mha benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full cutlass_attn --nheads="$nheads" \
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv cutlass_attn --nheads="$nheads" \
         --batch_size=$batch_size \
         --head_size=$head_size \
         --head_size_v=$head_size \
         --seq_length="$length" \
         --seq_length_kv=$length \
-        --causal=false  > $log_dir/cutlass_attention_ncu.csv
+        --causal=false  > $log_dir/cutlass_attention_ncu.csv
diff --git a/artifacts/run_all_ncu_flash2.sh b/artifacts/run_all_ncu_flash2.sh
@@ -9,5 +9,6 @@ mha_dir="$benchmark_dir/multi-head_attention/baseline"
 # 1. ncu test the mha benchmark
 echo "NCU profiling mha benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full python3 $mha_dir/test_pt_model.py > $log_dir/flash2_attention_ncu.csv
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv python3 $mha_dir/test_pt_model.py > $log_dir/flash2_attention_ncu.csv
 
diff --git a/artifacts/run_all_ncu_ft.sh b/artifacts/run_all_ncu_ft.sh
@@ -10,10 +10,12 @@ echo "NCU profiling mha benchmark"
 mha_dir="$benchmark_dir/multi-head_attention/fractaltensor/build"
 mha_exe="$mha_dir/main"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full $mha_exe > $log_dir/ft_attention_ncu.csv
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv $mha_exe > $log_dir/ft_attention_ncu.csv
 
 # 2. ncu test the bigbird benchmark
 bigbird_dir="$benchmark_dir/blocked_sparse_attention/fractaltensor/build"
 echo "NCU profiling BigBird benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full $bigbird_dir/bigbird > $log_dir/ft_bigbird_ncu.csv
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv $bigbird_dir/bigbird > $log_dir/ft_bigbird_ncu.csv
diff --git a/artifacts/run_all_ncu_pt.sh b/artifacts/run_all_ncu_pt.sh
@@ -10,4 +10,5 @@ bigbird_dir="$benchmark_dir/blocked_sparse_attention/pytorch"
 # 2. ncu test the bigbird benchmark
 echo "NCU profiling BigBird benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full python3 $bigbird_dir/main.py > $log_dir/pt_bigbird_ncu.csv
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv python3 $bigbird_dir/main.py > $log_dir/pt_bigbird_ncu.csv
diff --git a/artifacts/run_all_ncu_triton.sh b/artifacts/run_all_ncu_triton.sh
@@ -9,12 +9,14 @@ mha_dir="$benchmark_dir/multi-head_attention/baseline"
 # 1. ncu test the mha benchmark
 echo "NCU profiling mha benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full python3 $mha_dir/test_triton_model.py > $log_dir/triton_attention_ncu.csv
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv python3 $mha_dir/test_triton_model.py > $log_dir/triton_attention_ncu.csv
 
 
 bigbird_dir="$benchmark_dir/blocked_sparse_attention/triton"
 # 2. ncu test the bigbird benchmark
 echo "NCU profiling BigBird benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full python3 $bigbird_dir/main.py \
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv python3 $bigbird_dir/main.py \
     --default_test True > $log_dir/triton_bigbird_ncu.csv