From ee86969f6cd9ae7b7edbc4b4f9f6286a871e4acd Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 4 Jun 2026 03:57:59 +0000 Subject: [PATCH] Fix CUDA stream: use c10::cuda::getCurrentCUDAStream() directly in kernel launch --- dsv4/kernels/cuda/blackwell_swizzle.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dsv4/kernels/cuda/blackwell_swizzle.cu b/dsv4/kernels/cuda/blackwell_swizzle.cu index 71f2029b..1fb77914 100644 --- a/dsv4/kernels/cuda/blackwell_swizzle.cu +++ b/dsv4/kernels/cuda/blackwell_swizzle.cu @@ -104,12 +104,12 @@ void launch_blackwell_swizzle( // Pybind11 bindings for torch.utils.cpp_extension.load PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("blackwell_swizzle_32_4_4", [](at::Tensor input, at::Tensor output, int32_t rows, int32_t cols) { - auto stream = at::cuda::getCurrentCUDAStream().stream(); - launch_blackwell_swizzle( + auto stream = c10::cuda::getCurrentCUDAStream(); + blackwell_swizzle_32_4_4_kernel<<< + (rows * cols + 255) / 256, 256, 0, stream>>>( input.data_ptr(), output.data_ptr(), - rows, cols, - stream + rows, cols ); }, "Blackwell 32_4_4 scale swizzle"); }