Support Gumbel-Max Sampling Kernel #574
Conversation
STWMichael
changed the title
|
Please consider the latest version flashinfer-ai/flashinfer#2119 which significant improves performance. |
STWMichael
changed the title
STWMichael
changed the title
|
|
||
| template <uint32_t BLOCK_THREADS, | ||
| uint32_t VEC_SIZE, | ||
| int BATCH_SIZE, |
There was a problem hiding this comment.
Do we still need this template arg if we are passing in batch size as a function arg
| constexpr int CONSUMER_NUM_THREADS = 128; // Grace Hopper setting | ||
| #else | ||
| // Default settings for other architectures | ||
| constexpr int WORKER_NUM_THREADS = 128; |
There was a problem hiding this comment.
We also have WORKER_NUM_THREADS defined in the persistent_kernel.cuh file, not sure if it's a duplicated definition here.
JackFram
left a comment
JackFram
left a comment
There was a problem hiding this comment.
The sampling task implementation in this PR is already in a pretty good shape. I left some comments, and feel free to address them. Waiting for other reviewers input
| @@ -0,0 +1,228 @@ | |||
| /* Copyright 2023-2024 CMU | |||
| * | |||
There was a problem hiding this comment.
Per Apache 2.0 you have to keep the original copyright as well.
e.g.
Copyright 2023-2025 FlashInfer contributors
in addition to your own copyright.
There was a problem hiding this comment.
Sorry, I forgot to mention this. @STWMichae, please make sure to add the copyright before merging as well.
4 participants
Description of changes:
This PR adds support for Gumbel-Max Sampling Kernel in all architectures. The original kernel code is from flashinfer.
In order to support the sampling in MPK, we:
BLOCK_THREADSfrom 1024 to 128/256.batch_sizeand iterates through all batches sequentially.This PR also adds unit tests to verify the correctness of the kernel. It also provides a demo_sampling.py that shows the sampling effect during inference.
Related Issues:
#519
Linked Issues:
Issues closed by this PR: