[Kernel][RFC] Initial commit containing new Triton kernels for multi lora serving.

[FurtherAI]

SGMV Triton Kernels

Should Fix #2829, #4007, #4053, #4063, #3793, #4708

Modification of #5025 to change to contiguous weight format rather than S-LoRA paged weights.
In parallel with #5036

---

New Triton kernels for multi lora computation. These (should) handle any shape and data type, compute at the actual lora rank and also speed up for grouped lora requests (especially prefill).

• Replace Punica kernels.

ping @Yard1

[robertgshaw2-redhat]

@tlrmchlsmth

[tlrmchlsmth]

Just starting to take a look -- Could you explain a bit about the change to a contiguous weight format?

[FurtherAI]

@tlrmchlsmth So, S-LoRA puts the weights in a page that is [page_size, hidden_dim], each LoRA matrix has rank weights there. By contiguous, I mean going back to the format that is currently implemented for the Punica kernels, where the LoRA weights are stored in a stack that is padded to max rank (e.g. [max_loras, 1, hidden_dim, max_rank] for LoRA B)`.

@Yard1 requested this for simplicity

[Yard1]

This looks good - I think we can just merge it. Merging master should fix CI.

[FurtherAI]

@Yard1 You don't mean merge as in finish this PR already right? We still gotta remove punica

[Yard1]

@FurtherAI We can merge the kernels first, and then follow up with a PR to enable them. But either works!

[jeejeelee]

@Yard1 I'm still waiting for your decision, but you haven't provided any feedback, so I don't quite understand.

[jeejeelee]

@robertgshaw2-neuralmagic @tlrmchlsmth I would appreciate it if you could pay attention to my related PR #5036

[FurtherAI]

@Yard1 @tlrmchlsmth @robertgshaw2-neuralmagic
So, I am working on removing Punica and I think I've written all the code for it, working on testing now. Appreciate any pointers if anyone has any ideas about the errors I'm running into.

The first time I run a layer (currently testing column parallel in test_layers.py), it works. After that, it has a particular order of working and not working (seed is set to the same value at all times), being either off by 0.2 or 500 on average. This happens whether the layer is run in a loop or if the function is duplicated multiple times (they both repeat the same sequence). Pytest, Torch and Triton (autotuner disabled) all might persist something across different function calls, but I am baffled by this behavior, since I think they should be independent. Has anyone seen this before or have any ideas?

I pushed a commit with the current state.

Update

Fixed this, was due to using the tensor tracking lora ranks incorrectly. Triton was most likely accessing out of bounds but not throwing an error for it.

[FurtherAI]

@Yard1 Looks like I need Triton >= 2.2.0. Where all should the requirements be updated for this?
Also, specifically the Intel XPU install is failing. Naming issue or something, know what it is?

[chandan047]

Hi, our team is really looking forward to this merge! Is there an estimate on when this will be merged?

[FurtherAI]

@chandan047 This was largely taken over by #5036, though these kernels have the ability to compute at the actual rank of each LoRA, which I'm not sure was added in the other PR. They are also simpler, but I am not atm planning to try to merge this with the updated main. The other ones were merged into main so see if that works for your use case

[simon-mo]

Closed in favor of @jeejeelee's kernel. Thanks for the PR!