Automating warp-group partition

[htyu]

We are taking a simple heuristic alternative to user annotation to separate memory loads that are connected to a heavy dot operation into the producer group, while leaving everything else in the consumer groups. This appears to cover important kernels such as GEMM, FA, HSTU.

With this there is no user annotations required in the kernel code to enable WS. User just needs to specify some flags in the autotune configs to trigger WS, .e,g.,num_consumer_groups and num_buffers_warp_spec

    triton.Config(
        {
            "BLOCK_SIZE_M": 128,
            "BLOCK_SIZE_N": 256,
            "BLOCK_SIZE_K": 64,
            "GROUP_SIZE_M": 8,
            "NUM_CONSUMER_GROUPS": 2,
        },
        num_stages=2,
        num_warps=4,
        num_consumer_groups=2,
        num_buffers_warp_spec=3,
   ),

[manman-ren]

The rest looks good!

[manman-ren]

Curious on why we need to forward then backward, is it needed because of auto annotation?

[htyu]

Ideally backward and forward should be in a loop until no new ops gets a new task id. In practice, the backward->forward->backward pattern works well, similar to what has been done in data partition, because the forward pass is recursive.

Forward is not needed with user annotation as user will always annotate the final store.

[manman-ren]

should it be WarpGroupDot always?

[htyu]

You mean other mma operations? Ideally the heuristic should be able to handle all expansive ops, but for now we are limiting it to handling only wgmma, as that's the only warp-group level ops we have so far. Regular mma may not need a warp-group level of warp specialization.