ggml: CUMSUM and TRI (CPU, Metal, CUDA)

[gabe-l-hart]

Description

This PR builds on some of the work by @pwilkin in #16095 and extends the CPU implementations of CUMSUM and TRI to Metal and CUDA. It also extends type support to F16 and BF16.

The goal of this PR is to establish these two ops in the interest of both the DELTA_NET op for Qwen3-Next and the chunked implementation of the State Space Duality form of SSM_SCAN for faster prefill.

I'm putting this up for review now in case it helps with the Qwen3-Next work and to get feedback on the kernels. I'm quite novice at kernel development, so I suspect others may find significant optimizations for both Metal and CUDA.

[gabe-l-hart]

Yikes, looks like some alternate platforms that will need to be handled. I'll dig through some of these failures

[JohannesGaessler]

My opinion is that we should assert that the CPU implementation is correct before we move towards reviewing and merging this PR.

[pwilkin]

My opinion is that we should assert that the CPU implementation is correct before we move towards reviewing and merging this PR.

I think @gabe-l-hart moved the TRI and CUMSUM CPU implementation to this PR as well, so I guess it's a question of adding some testcases?

Those ops aren't very hard and as far as basic correctness goes I think I've verified them quite extensively during my fights with Qwen3 Next.

I've mimicked the basic logic for CUMSUM to be the same as SUM_ROWS, so it's basically always done on the first dimension.

[JohannesGaessler]

The thing is though that as of right now those ops aren't used anywhere on master. My opinion is that the new ops should be added in tandem with the model that needs them just in case it turns out that further changes are needed (even if it's unlikely).

[gabe-l-hart]

@JohannesGaessler That makes total sense. I'm continuing to work towards the SSD formulation for SSM_SCAN, so this PR is really just a checkpoint for the primitive ops. My goal is better performance for Granite 4 which is why I went through to Metal and CUDA here.

[giuseppe]

@JohannesGaessler That makes total sense. I'm continuing to work towards the SSD formulation for SSM_SCAN,

@gabe-l-hart FYI, I've worked on the implementation for SSM_SCAN for the Vulkan backend: #16463 and that indeed helped with Granite 4

[gabe-l-hart]

@giuseppe I saw that yesterday! That support will help a ton. The SSD formulation is an additional optimization on top of the recurrent formulation that should have big benefits for prefill with long context. It's mathematically equivalent, but much more efficient. The challenge I'm working through right now is how best to decompose the problem to minimize impact. One option is to write it as a sub-graph composed of smaller ops that is used when the sequence length is > 1, but this would have two problems:

1. It would break the ability to reuse the graph objects across generation steps

2. It would make backends that don't support the primitive ops (CUMSUM and TRI) fall back to CPU (exactly what you're trying to fix in your PR for SSM_SCAN)

The alternative is to implement it inside a single backend's SSM_SCAN. This has the advantage of being self-contained so it can be done on a per-backend basis, but it has the inverse problem of requiring each backend to implement it separately in order to get the performance boost. It's also much harder to write as a single kernel since it involves allocating temporary tensors of different size than the input or output tensors.