diff --git a/mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp b/mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp index cacfdf4ba5d6..f1a7969aacf0 100644 --- a/mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp +++ b/mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp @@ -1439,6 +1439,9 @@ getTuningProblemStr(RockGemmGemmWrapperInterface gemmGemmOp, problemOS << "false" << sep; problemOS << "-split_kv " << attentionOp.getSplitKV() << sep; + 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; diff --git a/mlir/test/fusion/mixr-attention-sliding-window-problem-key.mlir b/mlir/test/fusion/mixr-attention-sliding-window-problem-key.mlir new file mode 100644 index 000000000000..f5a0a0e01d25 --- /dev/null +++ b/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> + } +} diff --git a/mlir/test/perf-scripts/attention-tuning-db-compat.py b/mlir/test/perf-scripts/attention-tuning-db-compat.py index dabcb88c7756..7cd7224f8a8f 100644 --- a/mlir/test/perf-scripts/attention-tuning-db-compat.py +++ b/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,7 +104,7 @@ 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): @@ -111,7 +112,7 @@ def test_read_tuning_db_matches_pre_transbias_scale_bias_key(self): 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) diff --git a/mlir/test/rocmlir-gen/problem-key.mlir b/mlir/test/rocmlir-gen/problem-key.mlir index f1d0c690088a..d0068d4d742e 100644 --- a/mlir/test/rocmlir-gen/problem-key.mlir +++ b/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 diff --git a/mlir/utils/performance/analysis/quickTuningGen.py b/mlir/utils/performance/analysis/quickTuningGen.py index ae3b4c143e11..79a166515e4f 100644 --- a/mlir/utils/performance/analysis/quickTuningGen.py +++ b/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) diff --git a/mlir/utils/performance/attentionSweeps.py b/mlir/utils/performance/attentionSweeps.py index 998b589458aa..b43c87697fe7 100755 --- a/mlir/utils/performance/attentionSweeps.py +++ b/mlir/utils/performance/attentionSweeps.py @@ -115,7 +115,7 @@ def to_attn_config(params, options: Options) -> AttentionConfiguration: """Converts a sampled parameter tuple into a AttentionConfiguration instance.""" shape, perf = params - *shape_params, current_seqlen = shape + *shape_params, current_seqlen, sliding_window_size = shape dtype, g, slq, slk, nhq, nhkv, hdqk, hdv, scale, bias, tq, tk, tv, to, causal, rlse, split_kv = shape_params perf_str = f"attn:v3:{','.join(str(x) for x in perf)}" attn_config = AttentionConfiguration(dtype=dtype, @@ -138,8 +138,9 @@ def to_attn_config(params, options: Options) -> AttentionConfiguration: arch=options.arch, num_cu=options.num_cu, num_chiplets=options.num_chiplets, - perf_config=perf_str) - attn_config.current_seqlen = current_seqlen + perf_config=perf_str, + current_seqlen=current_seqlen, + sliding_window_size=sliding_window_size) return attn_config @@ -230,6 +231,10 @@ def sample_attn_shape(): seqlen_q = 1 if use_kvcache else random.randint(1, max_valid_seqlen) # SEQ_LEN_Q current_seqlen = gen_current_seqlens(g, seqlen_k) if use_kvcache else None + # Sliding-window masking is only valid in KV-cache mode. Retain plain + # KV-cache samples as well so both paths remain covered. + sliding_window_size = (random.randint(1, seqlen_k) + if use_kvcache and random.choice(BOOLS) else 0) num_heads_q = 1 num_heads_kv = 1 @@ -274,7 +279,8 @@ def sample_attn_shape(): random.choice(BOOLS), # causal return_lse, split_kv, - current_seqlen) + current_seqlen, + sliding_window_size) def _infer_instruction_set(arch: str, requested: str) -> str: @@ -327,7 +333,7 @@ def sample_attention_case(instruction_set: str, flags: list[str]): def _estimate_splitkv_extra_bytes(shape_sample: tuple) -> Optional[int]: - dtype, g, seq_len_q, _seq_len_k, num_heads_q, _num_heads_kv, _head_dim_qk, head_dim_v, _scale, _bias, _tq, _tk, _tv, _to, _causal, return_lse, split_kv, _current_seqlen = shape_sample + dtype, g, seq_len_q, _seq_len_k, num_heads_q, _num_heads_kv, _head_dim_qk, head_dim_v, _scale, _bias, _tq, _tk, _tv, _to, _causal, return_lse, split_kv, _current_seqlen, _sliding_window_size = shape_sample if split_kv <= 1: return 0 diff --git a/mlir/utils/performance/parameterSweeps.py b/mlir/utils/performance/parameterSweeps.py index 8799d579e968..310c2a321bc3 100755 --- a/mlir/utils/performance/parameterSweeps.py +++ b/mlir/utils/performance/parameterSweeps.py @@ -301,9 +301,6 @@ async def test_config(config, options: Options, paths: Paths) -> TestResult: else: rocmlir_gen_opts = config.generate_mlir_driver_commandline(' '.join(options.flags), kernel_repeats=None).split() - if getattr(config, "current_seqlen") is not None: - rocmlir_gen_opts.append( - f"--current_seq_len={','.join(map(str, config.current_seqlen))}") rocmlir_gen_opts.append('-pv') if (isinstance(config, perfRunner.AttentionConfiguration) and diff --git a/mlir/utils/performance/perfRunner.py b/mlir/utils/performance/perfRunner.py index cc4b990e8ce6..060f87b33480 100644 --- a/mlir/utils/performance/perfRunner.py +++ b/mlir/utils/performance/perfRunner.py @@ -17,7 +17,7 @@ import re from dataclasses import dataclass -from typing import Optional, Dict, Tuple +from typing import Optional, Dict, Tuple, List import numpy as np import pandas as pd from hip import hip @@ -821,8 +821,8 @@ def from_command_line(cls, argv, arch, num_cu, num_chiplets): def to_command_line(self): return ( - f"conv{ {'f32':'', 'f16':'fp16', 'bf16':'bfp16', 'i8':'int8','fp8_fp8':'fp8_fp8', 'fp8': 'fp8'}[self.datatype]} " - + f"-F { {'fwd':1, 'bwd':2, 'wrw':4}[self.direction]} " + + f"conv{dict(f32='', f16='fp16', bf16='bfp16', i8='int8', fp8_fp8='fp8_fp8', fp8='fp8')[self.datatype]} " + + f"-F {dict(fwd=1, bwd=2, wrw=4)[self.direction]} " + f"-f {inverse_filter_layouts(self.filter_layout)} -I {self.input_layout.upper()} " + f"-O {inverse_output_layouts(self.output_layout)} " + f"-n {self.n} -c {self.c} -H {self.hi} -W {self.wi} -k {self.k} " + @@ -1720,7 +1720,9 @@ def __init__(self, num_cu: int, num_chiplets: int, perf_config: str = '', - trans_bias: bool = False): + trans_bias: bool = False, + current_seqlen: Optional[List[int]] = None, + sliding_window_size: int = 0): if DATA_TYPES_ATTENTION is None: initialize_dtypes_attn() if dtype not in DATA_TYPES_ATTENTION: @@ -1746,6 +1748,10 @@ def __init__(self, self.causal = causal self.return_lse = return_lse self.split_kv = split_kv + # The window size changes the generated kernel and belongs in its + # tuning identity. Runtime sequence positions do not. + self.sliding_window_size = sliding_window_size + self.current_seqlen = current_seqlen self.arch = arch self.chip = GFX_CHIP_RE.search(arch).group(0) @@ -1782,9 +1788,9 @@ def table_entry(self, nanoseconds): values = [ self.datatype, self.chip, self.num_cu, self.num_chiplets, self.trans_q, self.trans_k, self.trans_v, self.trans_o, self.causal, self.return_lse, self.split_kv, - self.with_attn_scale, self.with_attn_bias, self.trans_bias, self.g, self.seq_len_q, - self.seq_len_k, self.num_heads_q, self.num_heads_kv, self.head_dim_qk, self.head_dim_v, - self.perfconfig, + self.sliding_window_size, self.with_attn_scale, self.with_attn_bias, self.trans_bias, + self.g, self.seq_len_q, self.seq_len_k, self.num_heads_q, self.num_heads_kv, + self.head_dim_qk, self.head_dim_v, self.perfconfig, self.compute_tflops(nanoseconds) ] assert (len(self.TABLE_COLUMNS) == len(values)) @@ -1810,7 +1816,10 @@ def generate_mlir_driver_commandline(self, rocmlir_gen_flags, kernel_repeats=MLI f"-with-attn-bias={self.with_attn_bias}", f"-transBias={self.trans_bias}", f"-transQ={self.trans_q}", f"-transK={self.trans_k}", f"-transV={self.trans_v}", f"-transO={self.trans_o}", f"-causal={self.causal}", f"-return_lse={self.return_lse}", - f"-split_kv={self.split_kv}", + f"-split_kv={self.split_kv}", *([f"-sliding_window_size={self.sliding_window_size}"] + if self.sliding_window_size > 0 else []), + *([f"-current_seq_len={','.join(map(str, self.current_seqlen))}"] + if self.current_seqlen else []), *(['--kernel-repeats', str(kernel_repeats)] if kernel_repeats is not None else []), f"--perf_config={self.perfconfig}" ]) @@ -1838,6 +1847,8 @@ def from_command_line(cls, argv, arch, num_cu, num_chiplets): causal = False return_lse = False split_kv = 1 + sliding_window_size = 0 + current_seqlen = None with_attn_scale = False with_attn_bias = False trans_bias = False @@ -1881,6 +1892,10 @@ def from_command_line(cls, argv, arch, num_cu, num_chiplets): return_lse = (val.lower() in ["1", "true"]) elif opt.endswith("-split_kv"): split_kv = int(val) + elif opt.endswith("-sliding_window_size"): + sliding_window_size = int(val) + elif opt.endswith("-current_seq_len"): + current_seqlen = [int(x) for x in val.split(",")] elif opt.endswith("-perf_config"): perf_config = val else: @@ -1914,7 +1929,9 @@ def from_command_line(cls, argv, arch, num_cu, num_chiplets): num_cu, num_chiplets, perf_config, - trans_bias=trans_bias) + trans_bias=trans_bias, + current_seqlen=current_seqlen, + sliding_window_size=sliding_window_size) def to_command_line(self): return ( @@ -1923,7 +1940,8 @@ def to_command_line(self): f"-transV {str(self.trans_v).lower()} -transO {str(self.trans_o).lower()} " + f"-causal {str(self.causal).lower()} " + f"-return_lse {str(self.return_lse).lower()} " + f"-split_kv {str(self.split_kv)} " + - f"-g {self.g} " + + (f"-sliding_window_size {str(self.sliding_window_size)} " + if self.sliding_window_size > 0 else "") + f"-g {self.g} " + f"-seq_len_q {str(self.seq_len_q)} -seq_len_k {str(self.seq_len_k)} -num_heads_q {str(self.num_heads_q)} -num_heads_kv {str(self.num_heads_kv)} -head_dim_qk {str(self.head_dim_qk)} -head_dim_v {str(self.head_dim_v)} " + f"-with-attn-scale {str(self.with_attn_scale).lower()} " + f"-with-attn-bias {str(self.with_attn_bias).lower()} " + diff --git a/mlir/utils/performance/reportUtils.py b/mlir/utils/performance/reportUtils.py index 387f86787b77..8f8a6525cc19 100644 --- a/mlir/utils/performance/reportUtils.py +++ b/mlir/utils/performance/reportUtils.py @@ -39,8 +39,8 @@ ] ATTN_TEST_PARAMETERS = [ 'DataType', 'Chip', 'numCU', 'numChiplets', 'TransQ', 'TransK', 'TransV', 'TransO', 'Causal', - 'ReturnLSE', 'SplitKV', 'WithAttnScale', 'WithAttnBias', 'TransBias', 'G', 'SeqLenQ', 'SeqLenK', - 'NumHeadsQ', 'NumHeadsKV', 'HeadDimQK', 'HeadDimV', 'PerfConfig' + 'ReturnLSE', 'SplitKV', 'SlidingWindowSize', 'WithAttnScale', 'WithAttnBias', 'TransBias', 'G', + 'SeqLenQ', 'SeqLenK', 'NumHeadsQ', 'NumHeadsKV', 'HeadDimQK', 'HeadDimV', 'PerfConfig' ] GEMM_GEMM_TEST_PARAMETERS = [ 'DataType', 'Chip', 'numCU', 'numChiplets', 'TransA', 'TransB', 'TransC', 'TransO', 'G', 'M',