SOLVE_TRI extension to more dimensions

[pwilkin]

Extend SOLVE_TRI with a tiled kernel to work on all dimensions. Asked my LLM friends to generate a tiled kernel, added some backend tests for it. Not really too worried about optimizations for now since the base case is 64, but adding this in case we want to experiment with different / configurable chunk size for hybrid models with chunking (like Qwen3Next or the hopefully-coming Kimi Linear).

[am17an]

Does this supersede #17703?

[pwilkin]

@am17an nope, that other one is optimization for the small kernel.

[am17an]

Does this kernel have similar performance to the templated one? If there is a not a huge performance difference it is much preferred to keep one kernel for future maintenance + binary size (a big difference in e2e performance of qwen3-next would justify its existence however)

[am17an]

This would probably stuff from the same bank conflict problem right?

[am17an]

const wherever

[JohannesGaessler]

Copypasting and modifying the code like this is not acceptable in terms of mainability. Instead modify the existing kernel to handle tiling for the case n_template == 0 && k_template == 0.

[JohannesGaessler]

I think you added this unintentionally, please remove it from this PR.

[pwilkin]

It adds itself automatically due to a GitHub hook installed by the server development code 😐

[JohannesGaessler]

Also be aware that this PR will need to be rebased on top of master.

[pwilkin]

I have a reworked approach using cuBlas trsmBatched for the big sizes instead, but I haven't been able to figure out some weird race conditions...

[pwilkin]

Okay, @JohannesGaessler @am17an question:

I rewrote the entire kernel to just use cublasStrsmBatched. For the small case, it is faster than my previous kernel but slower than the new kernel in master, but in general the numbers are like this:

Optimized kernel in master:

SOLVE_TRI(type=f32,ne_lhs=[64,64,4,4],ne_rhs=[32,64,4,4]):                   49140 runs -    21.26 us/run -   2.13 MFLOP/run - 100.18 GFLOPS

cuBLAS:

SOLVE_TRI(type=f32,ne_lhs=[64,64,4,4],ne_rhs=[32,64,4,4]):                   40950 runs -    28.86 us/run -   2.13 MFLOP/run -  73.80 GFLOPS
SOLVE_TRI(type=f32,ne_lhs=[128,128,4,2],ne_rhs=[32,128,4,2]):                16380 runs -    83.70 us/run -   4.23 MFLOP/run -  50.50 GFLOPS
SOLVE_TRI(type=f32,ne_lhs=[64,64,8,32],ne_rhs=[64,64,8,32]):                 10276 runs -   106.43 us/run -  68.16 MFLOP/run - 640.37 GFLOPS
SOLVE_TRI(type=f32,ne_lhs=[128,128,4,32],ne_rhs=[128,128,4,32]):              4070 runs -   255.32 us/run - 270.53 MFLOP/run -   1.06 TFLOPS
SOLVE_TRI(type=f32,ne_lhs=[256,256,4,2],ne_rhs=[128,256,4,2]):                7425 runs -   151.09 us/run -  67.37 MFLOP/run - 445.91 GFLOPS

Does it make sense to keep the optimized kernel for the small cases with the ~25% performance increase or should I just replace the entire kernel with the cuBLAS version?

[JohannesGaessler]

Is there an equivalent solution for HIP?

[pwilkin]

@JohannesGaessler from what I've checked there's both hipblasStrsm and hipblasStrsmBatched and there should be mappings for both.

[JohannesGaessler]

My opinion is that both a manually written kernel and a library call would both be fine. I don't think that performance numbers for test-backend-ops matter for this consideration, what matters is whether there is a meaningful difference in end-to-end performance with any actual models. 25% of 1% of the total runtime would be negligible.

[pwilkin]

Okay, I merged both versions then, I think we can keep it.

SOLVE_TRI is pretty mission-critical for tuning the new generation of hybrid models with recurrent logic, since its efficiency determines the chunk size that can be used (and thus potentially the graph size). It's only used now in Qwen3Next but probably Kimi Linear is next in line.