[NVIDIA] Fix an issue to use current stream for the nvfp4 quant #13632
Conversation
Signed-off-by: kaixih <kaixih@nvidia.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run 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 either: Add 🚀 |
|
cc. @trevor-m |
|
Adding @mgoin @LucasWilkinson who reviewed the original PR. Some context, we noticed that we should have use the current stream instead of from a pool when setting the cuda stream. This PR fixes it. PTAL. |
| auto stream = at::cuda::getStreamFromPool(false, input.get_device()); | ||
| if (stream == nullptr) { | ||
| std::cerr << "Warning: Null CUDA stream" << std::endl; | ||
| } | ||
| auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); |
There was a problem hiding this comment.
This looks good to me as it's more inline with other uses. What was the issue that popped up because of this?
There was a problem hiding this comment.
getStreamFromPool was giving a different stream for this kernel, so ops later on which are running on the main/current stream might try to access the result tensors before this kernel is finished. This was causing incorrect results and NaNs.
There was a problem hiding this comment.
We got some NaN values when the fp4 quant is issued to the wrong stream. I guess when we test e2e model. @trevor-m can provide more info.
…-project#13632) Signed-off-by: Michael Glass <mrglass@us.ibm.com>
This uses the current stream (as a common practice for other cuda ops).