Validity of Split-K, Revisited

[sohnryang]

Thanks for the question.

#63 was about pointwise -> (split-k) matmul. However, the question here is about (split-k) matmul -> pointwise. These two are completely different situations. As for when split-k can be applied, please use linear algebra to decide.

Re the validator, like I explained before, may rule out innovative scheduling. An example is the famous flash attention, which was not in anyone's "official" scheduling playbook until it's discovered.

Originally posted by @yarongmu-google in #67

I don't think this clarification is sufficient. As far as I know, linear algebra doesn't precisely specify what is valid for split-k execution (as it's an implementation detail rather than math), let alone dictate what kind of pointwise and matmul fusion is possible. For example, "using linear algebra to decide" is insufficient to rule out the option 1 (recomputation) in #63's point 2.

[natetyoung]

Additionally, there seems to be some confusion: my original question #63 was about pointwise -> (split-k,split-m,split-n) matmul, but the answer stated that the fusion was invalid because the pointwise could not participate in the split-k iteration. I don't see how this is the case; a pointwise whose output is the first input of a matmul operates on the m and k dimensions of that matmul, and so could participate in split-k. Is the issue that it couldn't participate in the split-n dimension?

[papp-pal-andras]

I agree that the answer to #63 was very surprising in several ways, or at least seems to contradict the spirit of "not ruling out innovative scheduling" discussed in other answers.

1. To my understanding, a split-k only affects the execution of the sink operator(s), right? Their input operators in the subgraph do not use split-k anymore, since then there would be no "unified execution grid", we'd have different number of iterations for different operators. See Example 5 startegy B, where Op[1] uses split-k tiling, but its input Op[0] does not use split-k tiling over Tensor[0] and Tensor[1].
   The example in issue Subgraph execution clarifications #63 seems identical to Example 5 strategy B, it's just that Op[0] is pointwise now. The answer to Subgraph execution clarifications #63 says that the "The entire subgraph executes at a single granularity", but this is as true for a Pointwise Op[0] as the MatMul Op[0] in Example 5: neither of them can participate in a split-k directly because they are not sink operations, but both can be executed in the same way as per linear algebra rules.

2. If you maintain that the answer to Subgraph execution clarifications #63 is correct, would the subgraph still be invalid with spatial tile size 256x256, i.e. if there were only 2 iterations, just according to the split-k dimension? There's no repetition of tiles in Tensor0 (the pointwise's input) now, but of course Op0 still cannot be split-k tiled since it is Pointwise.

3. Finally, a detour regarding the original question of Subgraph execution clarifications #63. If I understood the intent of the question right, this can occur also without any split-k or pointwise operators. In a simple chained MatMul like Example 5 (let's assume infinite fast memory, native granularity [64, 64]), with subgraph [0,1], execution granularity [64, 64, 128], traversal order [0,2,1,3], the situation is essentially the same: both row strips of Tensor[3] are computed twice via Op[0]. To my understanding, the answer to Subgraph execution clarifications #63 is then option 1, the computation happens twice. Is that correct?