[WIP] Working Grouped gemm with group ID #48
Conversation
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Co-authored-by: Lucas Wilkinson <[email protected]> Signed-off-by: ElizaWszola <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can do one of these:
🚀 |
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
…of tensors Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
| expert_offsets[:-1], problem_sizes2, | ||
| ab_strides2, ab_strides2, c_strides2) | ||
|
|
||
| return (c2[a_map.argsort()].view(m, topk, k) * |
There was a problem hiding this comment.
a_map = topk_ids.flatten().argsort()
...
a_map.argsort()
is there a way we can potentially compute these in fused way? (since the second argsort is computing the inverse of the first?) I feel like argsort can sometimes be nightmarishly slow
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
Signed-off-by: ElizaWszola <[email protected]>
| const torch::Tensor& topk_ids, torch::Tensor& expert_offsets, | ||
| torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2, | ||
| torch::Tensor& arg_sort, torch::Tensor& arg_sort_prim, | ||
| const int64_t num_experts, const int64_t n, const int64_t k) { | ||
| get_a_expert_offsets<<<1, num_experts>>>( |
There was a problem hiding this comment.
We should try to parallelize this across SMs, this will probably involve efficiently using atomicAdd, for the sorting portion I think we should be to do something like this:
https://github.com/vllm-project/vllm/blob/38094584566b89210a6f72a408eba1fae43c3d81/csrc/moe/moe_align_sum_kernels.cu#L260-L291
There was a problem hiding this comment.
For counting occurrences, it might make sense allocate k threads per expert (assert this is less than warpSize), and 256 / k experts (4 warps) per SM (~4-8 warps per SM probably makes sense, since each SM has 4 warp schedulers and on
an H100 we have 132 SMs so even with a low count of 4 experts per-SM we can pretty easily cover all models). Then this loop can be parallerized across the threads allocated to an expert:
for (int i = 0; i < topk_length; ++i) {
occurrences += (topk_ids[i] == expert_id);
}
and we can do an intra-warp reduction right at the end
No description provided.