Add sliding-window attention tuning support

mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp

@@ -1439,6 +1439,12 @@ getTuningProblemStr(RockGemmGemmWrapperInterface gemmGemmOp,
       problemOS << "false" << sep;
 
     problemOS << "-split_kv " << attentionOp.getSplitKV() << sep;
+    // Sliding-window size affects the generated kernel, but currentSeqLen is a
+    // runtime position and intentionally remains outside the tuning key.
+    // Omit the disabled value so existing attention keys retain their identity.
+    if (auto slidingWindowSize = attentionOp.getSlidingWindowSize();
+        slidingWindowSize && *slidingWindowSize > 0)
+      problemOS << "-sliding_window_size " << *slidingWindowSize << sep;
     problemOS << "-num_heads_q " << attentionOp.getNumHeadsQ() << sep;
     problemOS << "-num_heads_kv " << attentionOp.getNumHeadsKV() << sep;
     problemOS << "-g " << qShape[0] / attentionOp.getNumHeadsQ() << sep;

mlir/test/fusion/mixr-attention-sliding-window-problem-key.mlir

@@ -0,0 +1,60 @@
+// RUN: rocmlir-driver -kernel-pipeline=migraphx,highlevel %s | rocmlir-gen --emit-tuning-key - | FileCheck %s
+// CHECK: gfx942
+// CHECK-SAME: 304
+// CHECK-SAME: -t f16 -transQ false -transK true -transV false -transO false -causal false -return_lse false -split_kv 1 -sliding_window_size 3 -num_heads_q 1 -num_heads_kv 1 -g 2 -seq_len_q 1 -seq_len_k 8 -head_dim_qk 2 -head_dim_v 2 -with-attn-scale false -with-attn-bias false -transBias false
+// CHECK-NOT: current_seq_len
+
+module {
+  func.func @mlir_attention(%arg0: !migraphx.shaped<1x1xsi32, 1x1>,
+                            %arg1: !migraphx.shaped<1x6x1x2xf16, 12x2x2x1>,
+                            %arg2: !migraphx.shaped<1x2x8x2xf16, 32x16x2x1>,
+                            %arg3: !migraphx.shaped<1x2x8x2xf16, 32x16x2x1>) -> !migraphx.shaped<1x1x4xf16, 4x4x1> attributes {rock.kernel, rock.arch = "gfx942", rock.num_cu = 304 : i64} {
+    %0 = migraphx.literal(dense<-3> : tensor<1xsi32>) : <1xsi32, 1>
+    %1 = migraphx.literal(dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xsi32>) : <8xsi32, 1>
+    %2 = migraphx.literal(dense<4> : tensor<1x1xsi32>) : <1x1xsi32, 1x1>
+    %3 = migraphx.literal(dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xsi32>) : <8xsi32, 1>
+    %4 = migraphx.literal(dense<0xFC00> : tensor<1xf16>) : <1xf16, 1>
+    %5 = migraphx.literal(dense<5.000000e-01> : tensor<1xf16>) : <1xf16, 1>
+    %6 = migraphx.multibroadcast %1 {out_dyn_dims = [], out_lens = [1, 1, 1, 8]} : <8xsi32, 1> -> <1x1x1x8xsi32, 0x0x0x1>
+    %7 = migraphx.multibroadcast %0 {out_dyn_dims = [], out_lens = [1, 1, 1, 1]} : <1xsi32, 1> -> <1x1x1x1xsi32, 0x0x0x1>
+    %8 = migraphx.reshape %arg0 {dims = [1, 1, 1, 1]} : <1x1xsi32, 1x1> -> <1x1x1x1xsi32, 1x1x1x1>
+    %9 = migraphx.reshape %2 {dims = [1, 1, 1, 1]} : <1x1xsi32, 1x1> -> <1x1x1x1xsi32, 1x1x1x1>
+    %10 = migraphx.reshape %2 {dims = [1, 1, 1, 1]} : <1x1xsi32, 1x1> -> <1x1x1x1xsi32, 1x1x1x1>
+    %11 = migraphx.clip %8, %9, %10 : <1x1x1x1xsi32, 1x1x1x1>, <1x1x1x1xsi32, 1x1x1x1>, <1x1x1x1xsi32, 1x1x1x1> -> <1x1x1x1xsi32, 1x1x1x1>
+    %12 = migraphx.slice %arg1 {axes = [1], ends = [2], starts = [0]} : <1x6x1x2xf16, 12x2x2x1> -> <1x2x1x2xf16, 12x2x2x1>
+    %13 = migraphx.transpose %arg2 {permutation = [0, 1, 3, 2]} : <1x2x8x2xf16, 32x16x2x1> -> <1x2x2x8xf16, 32x16x1x2>
+    %14 = migraphx.dot %12, %13 : <1x2x1x2xf16, 12x2x2x1>, <1x2x2x8xf16, 32x16x1x2> -> <1x2x1x8xf16, 16x8x8x1>
+    %15 = migraphx.multibroadcast %4 {out_dyn_dims = [], out_lens = [1, 2, 1, 8]} : <1xf16, 1> -> <1x2x1x8xf16, 0x0x0x0>
+    %16 = migraphx.multibroadcast %5 {out_dyn_dims = [], out_lens = [1, 2, 1, 8]} : <1xf16, 1> -> <1x2x1x8xf16, 0x0x0x0>
+    %17 = migraphx.mul %14, %16 : <1x2x1x8xf16, 16x8x8x1>, <1x2x1x8xf16, 0x0x0x0> -> <1x2x1x8xf16, 16x8x8x1>
+    %18 = migraphx.add %11, %7 : <1x1x1x1xsi32, 1x1x1x1>, <1x1x1x1xsi32, 0x0x0x1> -> <1x1x1x1xsi32, 1x1x1x1>
+    %19 = migraphx.multibroadcast %18 {out_dyn_dims = [], out_lens = [8, 1, 1, 1]} : <1x1x1x1xsi32, 1x1x1x1> -> <8x1x1x1xsi32, 0x1x1x1>
+    %20 = migraphx.reshape %19 {dims = [8]} : <8x1x1x1xsi32, 0x1x1x1> -> <8xsi32, 0>
+    %21 = migraphx.greater %20, %3 : <8xsi32, 0>, <8xsi32, 1> -> <8xsi32, 1>
+    %22 = migraphx.convert %21 {target_type = 0 : i64} : <8xsi32, 1> to <8xsi8, 1>
+    %23 = migraphx.broadcast %22 {axis = 3 : i64, out_lens = [1, 2, 1, 8]} : <8xsi8, 1> -> <1x2x1x8xsi8, 0x0x0x1>
+    %24 = migraphx.where %23, %15, %17 : <1x2x1x8xsi8, 0x0x0x1>, <1x2x1x8xf16, 0x0x0x0>, <1x2x1x8xf16, 16x8x8x1> -> <1x2x1x8xf16, 16x8x8x1>
+    %25 = migraphx.multibroadcast %11 {out_dyn_dims = [], out_lens = [1, 1, 1, 8]} : <1x1x1x1xsi32, 1x1x1x1> -> <1x1x1x8xsi32, 1x1x1x0>
+    %26 = migraphx.greater %6, %25 : <1x1x1x8xsi32, 0x0x0x1>, <1x1x1x8xsi32, 1x1x1x0> -> <1x1x1x8xsi32, 0x0x0x1>
+    %27 = migraphx.convert %26 {target_type = 0 : i64} : <1x1x1x8xsi32, 0x0x0x1> to <1x1x1x8xsi8, 0x0x0x1>
+    %28 = migraphx.multibroadcast %27 {out_dyn_dims = [], out_lens = [1, 2, 1, 8]} : <1x1x1x8xsi8, 0x0x0x1> -> <1x2x1x8xsi8, 0x0x0x1>
+    %29 = migraphx.where %28, %15, %24 : <1x2x1x8xsi8, 0x0x0x1>, <1x2x1x8xf16, 0x0x0x0>, <1x2x1x8xf16, 16x8x8x1> -> <1x2x1x8xf16, 16x8x8x1>
+    %30 = migraphx.convert %29 {target_type = 2 : i64} : <1x2x1x8xf16, 16x8x8x1> to <1x2x1x8xf32, 16x8x8x1>
+    %31 = migraphx.reshape %30 {dims = [1, 2, 1, 8]} : <1x2x1x8xf32, 16x8x8x1> -> <1x2x1x8xf32, 16x8x8x1>
+    %32 = migraphx.reduce_max %31 {axes = [3]} : <1x2x1x8xf32, 16x8x8x1> -> <1x2x1x1xf32, 2x1x1x1>
+    %33 = migraphx.reshape %32 {dims = [1, 2, 1, 1]} : <1x2x1x1xf32, 2x1x1x1> -> <1x2x1x1xf32, 2x1x1x1>
+    %34 = migraphx.multibroadcast %33 {out_dyn_dims = [], out_lens = [1, 2, 1, 8]} : <1x2x1x1xf32, 2x1x1x1> -> <1x2x1x8xf32, 2x1x1x0>
+    %35 = migraphx.sub %30, %34 : <1x2x1x8xf32, 16x8x8x1>, <1x2x1x8xf32, 2x1x1x0> -> <1x2x1x8xf32, 16x8x8x1>
+    %36 = migraphx.exp %35 : <1x2x1x8xf32, 16x8x8x1> -> <1x2x1x8xf32, 16x8x8x1>
+    %37 = migraphx.reshape %36 {dims = [1, 2, 1, 8]} : <1x2x1x8xf32, 16x8x8x1> -> <1x2x1x8xf32, 16x8x8x1>
+    %38 = migraphx.reduce_sum %37 {axes = [3]} : <1x2x1x8xf32, 16x8x8x1> -> <1x2x1x1xf32, 2x1x1x1>
+    %39 = migraphx.reshape %38 {dims = [1, 2, 1, 1]} : <1x2x1x1xf32, 2x1x1x1> -> <1x2x1x1xf32, 2x1x1x1>
+    %40 = migraphx.multibroadcast %39 {out_dyn_dims = [], out_lens = [1, 2, 1, 8]} : <1x2x1x1xf32, 2x1x1x1> -> <1x2x1x8xf32, 2x1x1x0>
+    %41 = migraphx.div %36, %40 : <1x2x1x8xf32, 16x8x8x1>, <1x2x1x8xf32, 2x1x1x0> -> <1x2x1x8xf32, 16x8x8x1>
+    %42 = migraphx.convert %41 {target_type = 1 : i64} : <1x2x1x8xf32, 16x8x8x1> to <1x2x1x8xf16, 16x8x8x1>
+    %43 = migraphx.dot %42, %arg3 : <1x2x1x8xf16, 16x8x8x1>, <1x2x8x2xf16, 32x16x2x1> -> <1x2x1x2xf16, 4x2x2x1>
+    %44 = migraphx.transpose %43 {permutation = [0, 2, 1, 3]} : <1x2x1x2xf16, 4x2x2x1> -> <1x1x2x2xf16, 4x2x2x1>
+    %45 = migraphx.reshape %44 {dims = [1, 1, 4]} : <1x1x2x2xf16, 4x2x2x1> -> <1x1x4xf16, 4x4x1>
+    return %45 : !migraphx.shaped<1x1x4xf16, 4x4x1>
+  }
+}

mlir/test/perf-scripts/attention-tuning-db-compat.py

@@ -5,9 +5,10 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 """Pure-Python coverage for attention tuning DB compatibility.
 
-The tuning DB key for attention has grown optional boolean flags over time
-(``-with-attn-scale``, ``-with-attn-bias``, ``-transBias``). False-valued flags
-are identity cases, so old DB rows that omit them should still be readable:
+The tuning DB key for attention has grown optional fields over time
+(``-with-attn-scale``, ``-with-attn-bias``, ``-transBias``, and
+``-sliding_window_size``). Disabled fields are identity cases, so old DB rows
+that omit them should still be readable:
 ``read_tuning_db`` canonicalizes every stored key through
 ``AttentionConfiguration``, which re-adds the (default false) flags. True-valued
 flags describe different generated kernels and must not be silently matched
@@ -87,10 +88,10 @@ def setUp(self):
         """Give each test its own temporary file prefix."""
         self.tmp_prefix = Path(f"{TMP_PREFIX}.{self._testMethodName}")
 
-    def read_db_from_legacy_key(self, legacy_key):
-        """Read a tuning DB written with a legacy key (canonicalized on read)."""
+    def read_db_from_key(self, key):
+        """Read a tuning DB whose key is canonicalized on read."""
         path = Path(f"{self.tmp_prefix}.tsv")
-        write_tuning_db(path, legacy_key)
+        write_tuning_db(path, key)
         return read_tuning_db(str(path), AttentionConfiguration, NUM_CU, NUM_CHIPLETS)
 
     def db_key(self, config):
@@ -103,15 +104,15 @@ def test_read_tuning_db_matches_legacy_all_false_attention_flags(self):
         legacy_key = drop_flags(current_config.to_command_line(), " -with-attn-scale false",
                                 " -with-attn-bias false", " -transBias false")
 
-        db = self.read_db_from_legacy_key(legacy_key)
+        db = self.read_db_from_key(legacy_key)
         self.assertEqual(db.get(self.db_key(current_config)), PERFCONFIG)
 
     def test_read_tuning_db_matches_pre_transbias_scale_bias_key(self):
         """Scale+bias DB rows from before transBias still match non-transposed bias."""
         current_config = make_config("-with-attn-scale true -with-attn-bias true -transBias false")
         legacy_key = drop_flags(current_config.to_command_line(), " -transBias false")
 
-        db = self.read_db_from_legacy_key(legacy_key)
+        db = self.read_db_from_key(legacy_key)
         self.assertEqual(db.get(self.db_key(current_config)), PERFCONFIG)
 
     def test_read_tuning_db_keeps_true_scale_bias_distinct(self):
@@ -125,7 +126,7 @@ def test_read_tuning_db_keeps_true_scale_bias_distinct(self):
         scale_bias_config = make_config(
             "-with-attn-scale true -with-attn-bias true -transBias false")
 
-        db = self.read_db_from_legacy_key(legacy_all_false_key)
+        db = self.read_db_from_key(legacy_all_false_key)
         self.assertNotIn(self.db_key(scale_bias_config), db)
 
     def test_read_tuning_db_keeps_true_trans_bias_distinct(self):
@@ -140,11 +141,46 @@ def test_read_tuning_db_keeps_true_trans_bias_distinct(self):
             "-with-attn-scale false -with-attn-bias true -transBias true")
 
         self.assertIn("-transBias true", trans_bias_config.to_command_line())
-        db = self.read_db_from_legacy_key(legacy_all_false_key)
+        db = self.read_db_from_key(legacy_all_false_key)
         self.assertNotIn(self.db_key(trans_bias_config), db)
 
-    def test_quick_tuning_gen_defaults_missing_trans_bias_column(self):
-        """Legacy debug TSV rows without TransBias get a false default."""
+    def test_read_tuning_db_keeps_sliding_window_distinct(self):
+        """A sliding-window kernel must not match a row without its window."""
+        no_window_config = make_config(
+            "-with-attn-scale false -with-attn-bias false -transBias false")
+        sliding_window_config = make_config("-sliding_window_size 8 -with-attn-scale false "
+                                            "-with-attn-bias false -transBias false")
+
+        db = self.read_db_from_key(no_window_config.to_command_line())
+        self.assertNotIn(self.db_key(sliding_window_config), db)
+
+    def test_read_tuning_db_matches_pre_transbias_sliding_window_key(self):
+        """A sliding-window row from before transBias must still match."""
+        current_config = make_config("-sliding_window_size 8 -with-attn-scale false "
+                                     "-with-attn-bias false -transBias false")
+        legacy_key = drop_flags(current_config.to_command_line(), " -transBias false")
+
+        self.assertIn("-sliding_window_size 8", legacy_key)
+        db = self.read_db_from_key(legacy_key)
+        self.assertEqual(db.get(self.db_key(current_config)), PERFCONFIG)
+
+    def test_current_seq_len_is_runtime_only(self):
+        """Runtime positions reach rocmlir-gen without entering the tuning key."""
+        config = make_config("-sliding_window_size 8 -current_seq_len 4 "
+                             "-with-attn-scale false -with-attn-bias false -transBias false")
+
+        key = config.to_command_line()
+        driver_args = config.generate_mlir_driver_commandline("", kernel_repeats=None).split()
+        report_entry = config.table_entry(1.0)
+        self.assertIn("-sliding_window_size 8", key)
+        self.assertNotIn("current_seq_len", key)
+        self.assertIn("-sliding_window_size=8", driver_args)
+        self.assertEqual(driver_args.count("-current_seq_len=4"), 1)
+        self.assertEqual(report_entry["SlidingWindowSize"], 8)
+        self.assertNotIn("CurrentSeqLen", report_entry)
+
+    def test_quick_tuning_gen_defaults_missing_optional_columns(self):
+        """Legacy debug TSV rows get disabled optional-field defaults."""
         debug_path = Path(f"{self.tmp_prefix}.debug")
         debug_path.write_text(
             "DataType\tChip\tnumCU\tnumChiplets\tTransQ\tTransK\tTransV\tTransO\t"
@@ -154,17 +190,15 @@ def test_quick_tuning_gen_defaults_missing_trans_bias_column(self):
             f"False\tFalse\t1\tTrue\tTrue\t1\t16\t16\t1\t1\t32\t32\t{PERFCONFIG}\t1.0\n")
 
         df = load_data([str(debug_path)], no_splitk=False)
-        self.assertIn("TransBias", df.columns)
         self.assertTrue(df["TransBias"].eq(False).all())
+        self.assertTrue(df["SlidingWindowSize"].eq(0).all())
 
         grouped = df.groupby(get_target_columns("attention") + ["PerfConfig"],
                              as_index=False)["TFlops"].max()
         self.assertFalse(grouped.empty)
 
-    def test_quick_tuning_gen_fills_trans_bias_nan_in_mixed_files(self):
-        """Mixing legacy (no TransBias) and new (TransBias) TSVs must not drop
-        the legacy rows: concat leaves NaN in the legacy rows, and TransBias is
-        a groupby key, so an unfilled NaN would be silently dropped."""
+    def test_quick_tuning_gen_fills_optional_nan_in_mixed_files(self):
+        """Mixed legacy/current TSVs must retain rows with missing fields."""
         cols_no_tb = ("DataType\tChip\tnumCU\tnumChiplets\tTransQ\tTransK\tTransV\tTransO\t"
                       "Causal\tReturnLSE\tSplitKV\tWithAttnScale\tWithAttnBias\tG\tSeqLenQ\t"
                       "SeqLenK\tNumHeadsQ\tNumHeadsKV\tHeadDimQK\tHeadDimV\tPerfConfig\tTFlops\n")
@@ -173,23 +207,25 @@ def test_quick_tuning_gen_fills_trans_bias_nan_in_mixed_files(self):
             cols_no_tb + f"f16\tgfx950\t{NUM_CU}\t{NUM_CHIPLETS}\tFalse\tFalse\tFalse\tFalse\t"
             f"False\tFalse\t1\tTrue\tTrue\t1\t16\t16\t1\t1\t32\t32\t{PERFCONFIG}\t1.0\n")
 
-        cols_tb = (
+        cols_optional = (
             "DataType\tChip\tnumCU\tnumChiplets\tTransQ\tTransK\tTransV\tTransO\t"
-            "Causal\tReturnLSE\tSplitKV\tWithAttnScale\tWithAttnBias\tTransBias\tG\tSeqLenQ\t"
-            "SeqLenK\tNumHeadsQ\tNumHeadsKV\tHeadDimQK\tHeadDimV\tPerfConfig\tTFlops\n")
+            "Causal\tReturnLSE\tSplitKV\tSlidingWindowSize\tWithAttnScale\tWithAttnBias\t"
+            "TransBias\tG\tSeqLenQ\tSeqLenK\tNumHeadsQ\tNumHeadsKV\tHeadDimQK\tHeadDimV\t"
+            "PerfConfig\tTFlops\n")
         new_path = Path(f"{self.tmp_prefix}.new.debug")
-        new_path.write_text(
-            cols_tb + f"f16\tgfx950\t{NUM_CU}\t{NUM_CHIPLETS}\tFalse\tFalse\tFalse\tFalse\t"
-            f"False\tFalse\t1\tTrue\tTrue\tTrue\t1\t16\t16\t1\t1\t32\t32\t{PERFCONFIG}\t2.0\n")
+        new_path.write_text(cols_optional +
+                            f"f16\tgfx950\t{NUM_CU}\t{NUM_CHIPLETS}\tFalse\tFalse\tFalse\tFalse\t"
+                            f"False\tFalse\t1\t8\tTrue\tTrue\tTrue\t1\t16\t16\t1\t1\t32\t32\t"
+                            f"{PERFCONFIG}\t2.0\n")
 
         df = load_data([str(legacy_path), str(new_path)], no_splitk=False)
         self.assertFalse(df["TransBias"].isna().any())
+        self.assertFalse(df["SlidingWindowSize"].isna().any())
 
         grouped = df.groupby(get_target_columns("attention") + ["PerfConfig"],
                              as_index=False)["TFlops"].max()
-        # Both problems survive: legacy TransBias=False and new TransBias=True.
-        # Without the NaN fill, the legacy row would be dropped and only one
-        # group would remain.
+        # Both problems survive: legacy fields default to disabled, while the
+        # current row has TransBias=True and SlidingWindowSize=8.
         self.assertEqual(len(grouped), 2)
 
 

mlir/test/rocmlir-gen/problem-key.mlir

@@ -38,6 +38,16 @@
 // RUN: rocmlir-gen --arch gfx942 --operation attention -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -t i8 -g 8 | rocmlir-gen --emit-tuning-key - | FileCheck %s --check-prefixes=CHECK_I8_NO_SCALE_BIAS
 // CHECK_I8_NO_SCALE_BIAS: -t i8 {{.*}} -head_dim_v 32 -with-attn-scale false -with-attn-bias false -transBias false
 
+// Sliding-window size affects the generated kernel and is part of its tuning
+// identity. current_seq_len is runtime-only and is intentionally omitted.
+// RUN: rocmlir-gen --arch gfx942 --operation attention -current_seq_len=16 -sliding_window_size 8 -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -t f16 -g 1 | rocmlir-gen --emit-tuning-key - | FileCheck %s --check-prefixes=CHECK_SW
+// CHECK_SW: -t f16 -transQ false -transK false -transV false -transO false -causal false -return_lse false -split_kv 1 -sliding_window_size 8 -num_heads_q 1 -num_heads_kv 1 -g 1 -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -with-attn-scale false -with-attn-bias false -transBias false
+
+// Sliding-window and transposed-bias fields are independent and have stable
+// relative positions in the attention tuning key.
+// RUN: rocmlir-gen --arch gfx942 --operation attention -current_seq_len=16 -sliding_window_size 8 -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -t f16 -g 1 --with-attn-bias --transBias | rocmlir-gen --emit-tuning-key - | FileCheck %s --check-prefixes=CHECK_SW_TRANSBIAS
+// CHECK_SW_TRANSBIAS: -split_kv 1 -sliding_window_size 8 -num_heads_q 1 -num_heads_kv 1 -g 1 -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -with-attn-scale false -with-attn-bias true -transBias true
+
 // RUN: rocmlir-gen --arch gfx942 --operation conv -t f16 --fil_layout gkc01 --in_layout ngc01 --out_layout ngk01 --batchsize 64 --in_channels 256 --in_h 20 --in_w 20 --out_channels 256 --fil_h 7 --fil_w 7 --dilation_h 1 --dilation_w 1 --conv_stride_h 1 --conv_stride_w 1 --padding_h 3 --padding_w 3 --groupsize 256 --perf_config=v3:32,256,2,32,32,4,1,1,2,1,1 | rocmlir-gen --emit-tuning-key - | FileCheck %s  --check-prefixes=CHECK_DEPTHWISE_CONV
 // CHECK_DEPTHWISE_CONV: convfp16 -F 1 -f GNC01 -I NGC01 -O NGC01 -n 64 -c 256 -H 20 -W 20 -k 256 -y 7 -x 7 -p 3 -q 3 -u 1 -v 1 -l 1 -j 1 -g 256
 

mlir/utils/performance/analysis/quickTuningGen.py

@@ -23,9 +23,9 @@
     'DilationH', 'DilationW', 'StrideH', 'StrideW', 'PaddingH', 'PaddingW'
 ]
 ATTENTION_COLUMNS = [
-    'TransQ', 'TransK', 'TransV', 'TransO', 'Causal', 'ReturnLSE', 'SplitKV', 'WithAttnScale',
-    'WithAttnBias', 'TransBias', 'G', 'SeqLenQ', 'SeqLenK', 'NumHeadsQ', 'NumHeadsKV', 'HeadDimQK',
-    'HeadDimV'
+    'TransQ', 'TransK', 'TransV', 'TransO', 'Causal', 'ReturnLSE', 'SplitKV', 'SlidingWindowSize',
+    'WithAttnScale', 'WithAttnBias', 'TransBias', 'G', 'SeqLenQ', 'SeqLenK', 'NumHeadsQ',
+    'NumHeadsKV', 'HeadDimQK', 'HeadDimV'
 ]
 GEMM_GEMM_COLUMNS = ['TransA', 'TransB', 'TransC', 'TransO', 'G', 'M', 'K', 'N', 'O']
 CONV_GEMM_COLUMNS = [
@@ -197,6 +197,14 @@ def load_data(files, no_splitk):
         else:
             df['TransBias'] = df['TransBias'].fillna(False)
 
+    # Sliding windows are optional and were absent from legacy attention TSVs.
+    # Normalize both missing columns and NaNs introduced by mixed-file concat;
+    # otherwise groupby drops those legacy rows.
+    if 'SlidingWindowSize' not in df.columns:
+        df['SlidingWindowSize'] = 0
+    else:
+        df['SlidingWindowSize'] = df['SlidingWindowSize'].fillna(0)
+
     if no_splitk and not df.empty:
         # Filter out configs where Split-K != 1
         before = len(df)