This PR implements fused kernel for sampling functions (categorical sampling, top-k sampling, top-p sampling) of normalization functions (RMSNorm).
The top-k and top-p implementation uses reject sampling so that we don't need to use GPU sort. Our current implementation do not pin probability on shared memory, which could be further improved in the future.
This PR implements fused kernel for sampling functions (categorical sampling, top-k sampling, top-p sampling) of normalization functions (RMSNorm).
The top-k and top-p implementation uses reject sampling so that we don't need to use GPU sort. Our current implementation do not pin probability on shared memory, which could be further improved in the future.