Update attention sink examples to use 32 heads

- Modified the `heads` parameter in both `example_gqa_sink_fwd_bhsd_wgmma_pipelined.py` and `example_mha_sink_fwd_bhsd_wgmma_pipelined.py` from 1 to 32 to enhance performance in attention mechanisms.
- Ensured consistency across example scripts for improved usability and testing.

Author: LeiWang1999

examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py

@@ -366,7 +366,7 @@ def gen_inputs(B, H, Sq, Skv, D,
 
 def main(
     batch: int = 1,
-    heads: int = 1,
+    heads: int = 32,
     seq_q: int = 256,
     seq_kv: int = 256,
     dim: int = 128,

examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py

@@ -355,7 +355,7 @@ def gen_inputs(B, H, Sq, Skv, D) -> tuple[torch.Tensor, torch.Tensor, torch.Tens
 
 
 def main(batch: int = 1,
-         heads: int = 1,
+         heads: int = 32,
          seq_q: int = 256,
          seq_kv: int = 256,
          dim: int = 128,

src/op/builtin.cc

@@ -20,6 +20,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION(kDebugMergeSharedMemoryAllocations, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDisableTMALower, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDisableSafeMemoryLegalize, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDisableWarpSpecialized, Bool);
+TVM_REGISTER_PASS_CONFIG_OPTION(kDisableThreadStorageSync, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kConfigIndexBitwidth, Integer);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDisableDynamicTailSplit, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDynamicAlignment, Integer);