Update PyTorch pin (#4553)

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

.github/pins/pytorch.txt

@@ -1 +1 @@
-c2d1b225e62bdf3adbba91891f53bc60315adaac
+b5bfbba1841da810305262c7f47ee2dae54f335e

benchmarks/triton_kernels_benchmark/flex_attention_benchmark_causal_mask.py

@@ -13,30 +13,29 @@
 import torch._inductor.lowering
 import torch._inductor.kernel
 import torch._inductor.kernel.flex_attention as flex_attn
-import torch._inductor.virtualized
+from torch._inductor.template_heuristics import FlexConfig
 
 import triton_kernels_benchmark as benchmark_suit
 
 # Use TORCHINDUCTOR_MAX_AUTOTUNE_GEMM=1 or uncomment the following line to print the auto-tune results.
 # torch._inductor.config.max_autotune_gemm = True
 
 
-def get_xpu_config(*args, **kwargs):  # pylint: disable=unused-argument
-    # BLOCK_M, BLOCK_N, num_warps, num_stages
+def get_flex_attn_fwd_configs(*args, **kwargs):  # pylint: disable=unused-argument
     configs = [
-        (32, 16, 4, 2),
-        (128, 64, 16, 2),
-        (128, 64, 8, 2),
-        (128, 32, 16, 2),
-        (128, 32, 8, 2),
+        FlexConfig(32, 16, 2, 4),
+        FlexConfig(128, 64, 2, 16),
+        FlexConfig(128, 64, 2, 8),
+        FlexConfig(128, 32, 2, 16),
+        FlexConfig(128, 32, 2, 8),
     ]
     return configs
 
 
 # There is a auto-tuning requirement to get the best configuration for the flex attention.
 # The pytorch flex attention doesn't support auto-tuning by user by default.
-# Overriding the _get_xpu_config method to provide custom configurations for auto-tuning on XPU.
-flex_attn._get_xpu_config = get_xpu_config  # pylint: disable=protected-access
+# Overriding the get_flex_attention_fwd_configs method to provide custom configurations for auto-tuning on XPU.
+flex_attn.V.choices.get_flex_attention_fwd_configs = get_flex_attn_fwd_configs
 
 torch._dynamo.config.recompile_limit = 100  # pylint: disable=protected-access
 

scripts/patch-pytorch.sh

@@ -35,7 +35,7 @@ apply_patch() {
 echo "Applying PyTorch patches in $REPO_ROOT"
 
 # put your patch applies here
+# trigger build
 apply_patch ./patch/flex_attn_143553.patch
 apply_patch pytorch_fp64.patch
-apply_patch ./patch/Patch_torch_flex_attention_for_autotune_in_benchmark.patch
 apply_patch ./patch/pytorch_global_scratch.patch