Question about the cost function in persistent kernel scheduler

[Edenzzzz]

I noticed the current cost function models a linear computation load. Since attention has quadratic complexity, I tried using a quadratic workload estimation function as below, and benchmarked it, but it leads to lower throughput in all cases.
Could you help me understand why this may be the case? @yzh119
Thanks.

inline float cost_function(int qo_len, int kv_len) {
  return float(qo_len) * kv_len; // all kv are multipled by qo_len, then apply mask
    // if (kv_len >= qo_len) {
    //     // Right-aligned case: prefix + triangular
    //     int kv_init = kv_len - qo_len;
    //     float shared = float(qo_len) * float(kv_init);
    //     float tri    = float(qo_len) * float(qo_len + 1) / 2.0f;
    //     return shared + tri;
    // } else {
    //     // kv_len < qo_len: classic causal sum
    //     float tri = float(kv_len) * float(kv_len + 1) / 2.0f;
    //     float rem = float(qo_len - kv_len) * float(kv_len);
    //     return tri + rem;
    // }
}

[yzh119]

The qo_len here is basically a tile size which is usually a fixed value (16/128/64), so I guess whether using quadratic estimation or not doesn't make too much difference.

You probably shouldn't use logic such kv_len >= qo_len because here qo_len means tile size on query dimension, not the real query length (the naming is confusing and we should improve it).

[Edenzzzz]

Thanks for the clarification. Indeed it's a tile size, which is why I thought of qo_len * kv_len.
But with a small tile size the quadraticness will not dominate, and it seems hard to accurately model the compute load due to a number of other ops (online softmax, block reduce)