[NPU] Add softmax implementation

[lowdy1]

Summary

This PR adds a Softmax implementation for NPU. It includes a single-block forward kernel for smaller column sizes, as well as a multi-block kernel for large column sizes to avoid NPU UB overflow.

Testing Done

Test done with python -m pytest test/transformers/test_softmax.py
Hardware Type: Atlas 800I A2(32G)

• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[noemotiovon]

Some of my own thoughts — hope they’re helpful.

[noemotiovon]

In Triton kernel implementations on NPU, it’s best for the first dimension of the grid to be num_cores.

[noemotiovon]

ditto

[noemotiovon]

I don’t think we need two separate kernels. Also, the current approach of launching the kernel row by row is not optimal. It’s clear that the UB (unified buffer) isn’t being fully utilized.

I believe we should use a single unified kernel, with two nested loops inside: the outer loop iterates over rows, and the inner loop iterates over columns. If one kernel can process multiple rows, it can make better use of the UB.

However, this might introduce some additional overhead, so we still need to benchmark and see whether it actually improves performance.

[Tcc0403]

Agreed, multi-block kernel plus grid stride loop over rows by num_prorams should work.

[lowdy1]

Thanks for your suggestions.