]> git.djapps.eu Git - pkg/ggml/sources/whisper.cpp/commit
CUDA: add a fused top-K MoE kernel (llama/16130)
authorAman Gupta <redacted>
Thu, 25 Sep 2025 14:35:05 +0000 (22:35 +0800)
committerGeorgi Gerganov <redacted>
Mon, 29 Sep 2025 12:18:10 +0000 (15:18 +0300)
commitd9bf63cfb8ee1a44fa27547d95f59bae609209e5
treee99951e1ba648075271ebb32be7b56b1f530928a
parent24ea5476de65d5a6d716c7d8e775d6b6b037345a
CUDA: add a fused top-K MoE kernel (llama/16130)

* CUDA: add a fused top-K MoE kernel

This kernel does the following:
1. softmax over the logits per token [n_experts, n_tokens]
2. argmax reduce over the top-k (n_experts_used) logits
3. write weights + ids to global memory

It is intended as fusion of softmax->top-k->get_rows pipeline for MoE models

* Refactor into ggml_cuda_should_use_topk_moe

* Review: Use better coalescing pattern, use WARP_SIZE, store logits into registers before

* Review: format + micro-optimizations

* Fix bug: fix tie breakers

* Add optional norm + clean-up code

* Use smem for final write

* Add bounds check

* Use better memory pattern for writeback
ggml/src/ggml-cuda/ggml-cuda.cu
ggml/src/ggml-cuda/topk-moe.cu [new file with mode: 0644]
ggml/src/ggml-cuda/topk-moe.cuh [new file with mode: 0644]