align num_stages=3 between PoC and main (#2338)

Although `num_stages` in python file is always 3. We did not accept
`num_stages` as a parameter of `add_prefetch_block` pass. So PoC branch
uses `num_stages=3`, main branch uses `num_stages=4`.

Author: quintinwang5

benchmarks/triton_kernels_benchmark/flash_attention_fwd_benchmark.py

@@ -159,7 +159,7 @@ def forward(q, k, v, causal, sm_scale):
     o = torch.empty_like(q, dtype=torch.float32)
     BLOCK_M = 128
     BLOCK_N = 64 if Lk <= 64 else 32
-    num_stages = 4 if Lk <= 64 else 3
+    num_stages = 3
     num_warps = 8 if Lq == 64 else 16
     causal = False
     stage = 3 if causal else 1