A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 22:15 UTC
[NVIDIA] Fix an issue to use current stream for the nvfp4 quant (#13632)
Kaixi Hou committed
Feb 20, 2025 at 22:01 UTC
27a09dc52c8317b531b6d2b862198a8a0d2a88eb
1 file changed
+1
-4
csrc/quantization/fp4/nvfp4_quant_kernels.cu
+1
-4
@@ -348,10 +348,7 @@ void scaled_fp4_quant_sm100a(torch::Tensor const& output,
348
auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
349
auto output_ptr = static_cast<int64_t*>(output.data_ptr());
350
at::cuda::CUDAGuard device_guard{(char)input.get_device()};
351
- auto stream = at::cuda::getStreamFromPool(false, input.get_device());
352
- if (stream == nullptr) {
353
- std::cerr << "Warning: Null CUDA stream" << std::endl;
354
- }
351
+ auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
352
353
// We don't support e8m0 scales at this moment.
354
bool useUE8M0 = false;