[CK_TILE] Add sequence padding and variable length support in fmha (a… (#2851)

* [CK_TILE] Add sequence padding and variable length support in fmha (and v3)

 - Group Mode Padding: Introduces the `-s_qpad` argument to support
   physically padded layouts. Kernels now use padded start pointers
   (`seqstart_padded_*_ptr`) for memory addressing.

 - Batch Mode Variable Length: Adds `-q_eff_lens` and `-kv_eff_lens`
   arguments for efficient processing of variable-length sequences by
   passing cumulative effective lengths (`cu_seqlen_*_ptr`) to the kernel.

 - FMHA examples: Support padding and variable length both in
   group and batch mode. Dispatcher is updated as well (dispatch to
   kPadSeqLenK enabled pipeline).

 - New padding test cases: Add padding test cases to `smoke_test_fwd.sh`,
   and add benchmarks to `benchmark_fwd.sh` and `benchmark_fwd_v3.sh` as well.
   These test cases and benchmarks that specifically validate/benchmark the
   new padding and variable-length functionalities in both group and batch modes.

* [CK_TILE] Fix build error in fmha unit tests

---------

Co-authored-by: Po Yen Chen <[email protected]>
Co-authored-by: Yi DING <[email protected]>

Author: 3 people

example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py

@@ -259,11 +259,11 @@ def seqtune(self) -> str:
     def skcheck(self) -> str:
         if self.mode == 'group': return 'true/*group mode skpad always true*/'                  # group mode only generate spad/skpad == true
         if self.pipeline_tag == 'qr_async':
-            if self.skpad == 't' : return f'a.seqlen_k == 0 || a.seqlen_k % {self.bn0} != 0'
-            else :                 return f'a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0'
+            if self.skpad == 't' : return f'(a.cu_seqlen_kv_ptr != nullptr) || (a.seqlen_k == 0 || a.seqlen_k % {self.bn0} != 0)'
+            else :                 return f'(a.cu_seqlen_kv_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0)'
         elif self.pipeline_tag in ['qr', 'qs']:
             if self.skpad == 't' : return f'true /*a.seqlen_k % {self.bn0} != 0*/' # TODO: order of get_pipelines() matters! (ugly)
-            else :                return f'a.seqlen_k % {self.bn0} == 0'
+            else :                 return f'(a.cu_seqlen_kv_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0)'
         elif self.pipeline_tag == 'qr_async_trload':
             if self.skpad == 't' : return 'true'
             else:                  return 'true'

example/ck_tile/01_fmha/example_fmha_fwd.cpp

@@ -33,6 +33,10 @@ auto create_args(int argc, char* argv[])
                 "0",
                 "seqlen_k for new key/value, 0 means not to use this at all; "
                 "-1 to choose s_knew in [1, s] randomly.")
+        .insert("s_qpad",
+                "-1",
+                "seqlen_q stride between 2 batches (group-mode optional).\n"
+                "Provide positive strides per-batch to simulate physical padding on Q.")
         .insert("s_kpad",
                 "-1",
                 "seqlen_k stride between 2 batches, currently used in group-mode only\n"
@@ -107,7 +111,15 @@ auto create_args(int argc, char* argv[])
         .insert("warmup", "5", "number of iterations before benchmark the kernel")
         .insert("repeat", "20", "number of iterations to benchmark the kernel")
         .insert("json", "0", "0: No Json, 1: Dump Results in Json format")
-        .insert("jsonfile", "fmha_fwd.json", "json file name to dump results");
+        .insert("jsonfile", "fmha_fwd.json", "json file name to dump results")
+        .insert("q_eff_lens",
+                "",
+                "Batch-mode only: per-batch effective seqlen for Q (exclude PAD).\n"
+                "Comma-separated list of length 'b'. If empty, no override.")
+        .insert("kv_eff_lens",
+                "",
+                "Batch-mode only: per-batch effective seqlen for KV (exclude PAD).\n"
+                "Comma-separated list of length 'b'. If empty, no override.");
 
     bool result = arg_parser.parse(argc, argv);
     return std::make_tuple(result, arg_parser);
@@ -127,6 +139,9 @@ auto run(const ck_tile::ArgParser& arg_parser)
     ck_tile::index_t hdim_v          = arg_parser.get_int("d_v");
     ck_tile::index_t seqlen_knew     = arg_parser.get_int("s_knew");
     auto seqlen_kpads                = arg_parser.get_int_vec("s_kpad");
+    auto seqlen_qpads                = arg_parser.get_int_vec("s_qpad");
+    auto q_eff_lens_per_batch        = arg_parser.get_int_vec("q_eff_lens");
+    auto kv_eff_lens_per_batch       = arg_parser.get_int_vec("kv_eff_lens");
     ck_tile::index_t rotary_dim      = arg_parser.get_int("rotary_dim");
     bool i_perm                      = arg_parser.get_bool("iperm");
     bool o_perm                      = arg_parser.get_bool("operm");
@@ -174,7 +189,10 @@ auto run(const ck_tile::ArgParser& arg_parser)
                                         hdim_q,
                                         hdim_v,
                                         seqlen_knew,
+                                        seqlen_qpads,
                                         seqlen_kpads,
+                                        q_eff_lens_per_batch,
+                                        kv_eff_lens_per_batch,
                                         rotary_dim,
                                         i_perm,
                                         o_perm,

example/ck_tile/01_fmha/example_fmha_fwd_v3.cpp

@@ -52,7 +52,16 @@ auto parse_cmd_args(int argc, char* argv[]) -> std::pair<bool, ck_tile::ArgParse
                 "random seed used for initializing input tensors. 0 for "
                 "non-deterministic seed")
         .insert("warmup", "5", "number of iterations before benchmark the kernel")
-        .insert("repeat", "30", "number of iterations to benchmark the kernel");
+        .insert("repeat", "30", "number of iterations to benchmark the kernel")
+        // Optional effective seqlen override (exclude PAD) for batch mode
+        .insert("q_eff_lens",
+                "",
+                "Batch-mode only: per-batch effective seqlen for Q (exclude PAD).\n"
+                "Comma-separated list of length 'b'. If empty, no override.")
+        .insert("kv_eff_lens",
+                "",
+                "Batch-mode only: per-batch effective seqlen for KV (exclude PAD).\n"
+                "Comma-separated list of length 'b'. If empty, no override.");
 
     bool result = arg_parser.parse(argc, argv);
     return std::make_pair(result, arg_parser);
@@ -111,6 +120,8 @@ struct Problem
 
         input_layout  = args.get_int("iperm") == 1 ? TensorLayout::bhsd : TensorLayout::bshd;
         output_layout = args.get_int("operm") == 1 ? TensorLayout::bhsd : TensorLayout::bshd;
+        q_eff_lens    = args.get_int_vec("q_eff_lens");
+        kv_eff_lens   = args.get_int_vec("kv_eff_lens");
     }
 
     std::vector<ck_tile::index_t> get_query_shape() const
@@ -172,6 +183,8 @@ struct Problem
     mask_info mask;
     TensorLayout input_layout;
     TensorLayout output_layout;
+    std::vector<int> q_eff_lens;
+    std::vector<int> kv_eff_lens;
 };
 
 struct RunConfig
@@ -326,8 +339,10 @@ bool run_impl(const Problem& problem, const RunConfig& run_config)
     q_buf.ToDevice(q.data());
     k_buf.ToDevice(k.data());
     v_buf.ToDevice(v.data());
+    // Ensure output buffer is zero-initialized so padded regions compare cleanly
+    o_buf.SetZero();
 
-    ck_tile::fmha_fwd_v3_args args;
+    ck_tile::fmha_fwd_v3_args args{};
 
     args.data_type     = problem.data_type;
     args.batch         = problem.batch;
@@ -380,6 +395,60 @@ bool run_impl(const Problem& problem, const RunConfig& run_config)
                               : problem.seqlen_q * problem.hdim;
     args.batch_stride_o = problem.seqlen_q * problem.nhead_q * problem.hdim;
 
+    // Optional cumulative seqlen overrides (exclude PAD)
+    const bool has_varlen_q = !problem.q_eff_lens.empty() && problem.q_eff_lens[0] != -1;
+    const bool has_varlen_k = !problem.kv_eff_lens.empty() && problem.kv_eff_lens[0] != -1;
+
+    auto make_effective_vec = [&](const std::vector<int>& opt_vec, ck_tile::index_t fallback) {
+        std::vector<ck_tile::index_t> eff;
+        if(!opt_vec.empty() && opt_vec[0] != -1)
+        {
+            eff.assign(opt_vec.begin(), opt_vec.end());
+            if(eff.size() < static_cast<size_t>(problem.batch))
+            {
+                eff.resize(problem.batch, eff.back());
+            }
+        }
+        else
+        {
+            eff.assign(problem.batch, fallback);
+        }
+        return eff;
+    };
+
+    const auto eff_q_vec  = make_effective_vec(problem.q_eff_lens, problem.seqlen_q);
+    const auto eff_kv_vec = make_effective_vec(problem.kv_eff_lens, problem.seqlen_k);
+
+    // Calculate cumulative sums for kernel arguments if varlen is used
+    std::vector<ck_tile::index_t> cuq_cum, cukv_cum;
+    auto calculate_cumulative = [&](const std::vector<ck_tile::index_t>& per_batch_vec,
+                                    std::vector<ck_tile::index_t>& cum_vec) {
+        cum_vec.resize(per_batch_vec.size() + 1);
+        cum_vec[0] = 0;
+        for(std::size_t i = 0; i < per_batch_vec.size(); ++i)
+            cum_vec[i + 1] = cum_vec[i] + per_batch_vec[i];
+    };
+
+    if(has_varlen_q)
+    {
+        calculate_cumulative(eff_q_vec, cuq_cum);
+    }
+    if(has_varlen_k)
+    {
+        calculate_cumulative(eff_kv_vec, cukv_cum);
+    }
+
+    ck_tile::DeviceMem cuq_buf(!cuq_cum.empty() ? cuq_cum.size() * sizeof(ck_tile::index_t) : 0);
+    ck_tile::DeviceMem cukv_buf(!cukv_cum.empty() ? cukv_cum.size() * sizeof(ck_tile::index_t) : 0);
+    cuq_buf.ToDevice(!cuq_cum.empty() ? cuq_cum.data() : nullptr);
+    cukv_buf.ToDevice(!cukv_cum.empty() ? cukv_cum.data() : nullptr);
+    args.cu_seqlen_q_ptr =
+        !cuq_cum.empty() ? reinterpret_cast<const ck_tile::index_t*>(cuq_buf.GetDeviceBuffer())
+                         : nullptr;
+    args.cu_seqlen_kv_ptr =
+        !cukv_cum.empty() ? reinterpret_cast<const ck_tile::index_t*>(cukv_buf.GetDeviceBuffer())
+                          : nullptr;
+
     ck_tile::stream_config stream_config{nullptr,
                                          true,
                                          /*log_level=*/0,
@@ -442,15 +511,72 @@ bool run_impl(const Problem& problem, const RunConfig& run_config)
         o_ref = o_ref.transpose({0, 2, 1, 3});
     }
 
-    host::fmha_fwd<float, DataType>(q,
-                                    k,
-                                    v,
-                                    problem.mask,
-                                    o_ref,
-                                    ck_tile::identity{},
-                                    ck_tile::identity{},
-                                    ck_tile::identity{},
-                                    ck_tile::scales{problem.softmax_scale});
+    // If variable lengths are provided, compute per-batch references
+    // with the effective lengths; else compute a single full reference.
+    if(has_varlen_q || has_varlen_k)
+    {
+        // Variable-length aware verification: zero-fill padded region and only compute valid part.
+        o_ref.SetZero();
+
+        for(int b = 0; b < problem.batch; ++b)
+        {
+            const ck_tile::index_t seqlen_q_eff  = eff_q_vec[b];
+            const ck_tile::index_t seqlen_kv_eff = eff_kv_vec[b];
+
+            if(seqlen_q_eff <= 0 || seqlen_kv_eff <= 0)
+                continue;
+
+            // Slice current batch from inputs (bshd) and build single-batch tensors
+            ck_tile::HostTensor<DataType> q_b({1, seqlen_q_eff, problem.nhead_q, problem.hdim});
+            ck_tile::HostTensor<DataType> k_b({1, seqlen_kv_eff, problem.nhead_kv, problem.hdim});
+            ck_tile::HostTensor<DataType> v_b({1, seqlen_kv_eff, problem.nhead_kv, problem.hdim});
+            ck_tile::HostTensor<DataType> o_b({1, seqlen_q_eff, problem.nhead_q, problem.hdim});
+
+            // Copy effective region
+            q_b.ForEach([&](auto& self, auto idx) {
+                // idx: [0, s, h, d]
+                self(idx) = q(b, idx[1], idx[2], idx[3]);
+            });
+            k_b.ForEach([&](auto& self, auto idx) { self(idx) = k(b, idx[1], idx[2], idx[3]); });
+            v_b.ForEach([&](auto& self, auto idx) { self(idx) = v(b, idx[1], idx[2], idx[3]); });
+
+            // Compute reference for this batch segment (host::fmha_fwd expects bshd tensors)
+            host::fmha_fwd<float, DataType>(q_b,
+                                            k_b,
+                                            v_b,
+                                            problem.mask,
+                                            o_b,
+                                            ck_tile::identity{},
+                                            ck_tile::identity{},
+                                            ck_tile::identity{},
+                                            ck_tile::scales{problem.softmax_scale});
+
+            // Scatter into o_ref's bshd descriptor memory
+            for(int s = 0; s < seqlen_q_eff; ++s)
+            {
+                for(int h = 0; h < problem.nhead_q; ++h)
+                {
+                    for(int d = 0; d < problem.hdim; ++d)
+                    {
+                        o_ref(b, s, h, d) = o_b(0, s, h, d);
+                    }
+                }
+            }
+        }
+    }
+    else
+    {
+        // No varlen override: compute the full reference once
+        host::fmha_fwd<float, DataType>(q,
+                                        k,
+                                        v,
+                                        problem.mask,
+                                        o_ref,
+                                        ck_tile::identity{},
+                                        ck_tile::identity{},
+                                        ck_tile::identity{},
+                                        ck_tile::scales{problem.softmax_scale});
+    }
 
     ck_tile::HostTensor<DataType> o(problem.get_output_shape());
     o_buf.FromDevice(o.data());

example/ck_tile/01_fmha/fmha_fwd.hpp

@@ -162,11 +162,20 @@ struct fmha_fwd_args
     void* lse_ptr;
     void* o_ptr;
 
+    // Optional cumulative sequence length arrays
+    // Batch mode: cu_seqlen_* override effective per-batch lengths (exclude PAD)
+    const ck_tile::index_t* cu_seqlen_q_ptr  = nullptr; // [batch+1]
+    const ck_tile::index_t* cu_seqlen_kv_ptr = nullptr; // [batch+1]
+
     const void* seqstart_q_ptr;
     const void* seqstart_k_ptr;
     const void*
         seqlen_k_ptr; // only used if both 'seqstart_q_ptr' & 'seqstart_k_ptr' are not nullptr
 
+    // Group mode: seqstart_padded_* provide physical starts including PAD (optional)
+    const void* seqstart_padded_q_ptr = nullptr; // [batch+1]
+    const void* seqstart_padded_k_ptr = nullptr; // [batch+1]
+
     ck_tile::index_t seqlen_q;
     ck_tile::index_t seqlen_k;
     ck_tile::index_t batch;
@@ -554,7 +563,9 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args)
                                              args.min_seqlen_q,
                                              args.p_drop,
                                              args.s_randval,
-                                             args.drop_seed_offset);
+                                             args.drop_seed_offset,
+                                             args.seqstart_padded_q_ptr,
+                                             args.seqstart_padded_k_ptr);
         }
         else
         { // create batch mode kernel arguments
@@ -600,7 +611,9 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args)
                                              args.mask_type,
                                              args.p_drop,
                                              args.s_randval,
-                                             args.drop_seed_offset);
+                                             args.drop_seed_offset,
+                                             args.cu_seqlen_q_ptr,
+                                             args.cu_seqlen_kv_ptr);
         }
     }();