test fix

Author: LeiWang1999

3rdparty/tvm

@@ -366,9 +366,9 @@ def gen_inputs(B, H, Sq, Skv, D,
 
 def main(
     batch: int = 1,
-    heads: int = 64,
-    seq_q: int = 4096,
-    seq_kv: int = 4096,
+    heads: int = 1,
+    seq_q: int = 256,
+    seq_kv: int = 256,
     dim: int = 128,
     groups: int = 8,
     window_size: int | None = None,

examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py

@@ -229,10 +229,10 @@ def gen_inputs(B, H, Sq, Skv, D) -> tuple[torch.Tensor, torch.Tensor, torch.Tens
     return query, key, value, sinks
 
 
-def main(batch: int = 8,
-         heads: int = 32,
-         seq_q: int = 4096,
-         seq_kv: int = 4096,
+def main(batch: int = 1,
+         heads: int = 1,
+         seq_q: int = 256,
+         seq_kv: int = 256,
          dim: int = 128,
          window_size: int | None = None,
          tune: bool = False):

examples/attention_sink/example_mha_sink_fwd_bhsd.py

@@ -354,10 +354,10 @@ def gen_inputs(B, H, Sq, Skv, D) -> tuple[torch.Tensor, torch.Tensor, torch.Tens
     return query, key, value, sinks
 
 
-def main(batch: int = 8,
-         heads: int = 32,
-         seq_q: int = 4096,
-         seq_kv: int = 4096,
+def main(batch: int = 1,
+         heads: int = 1,
+         seq_q: int = 256,
+         seq_kv: int = 256,
          dim: int = 128,
          window_size: int | None = None,
          tune: bool = False):