Add CUDAGuard to ensure correct device (#5113)

Summary:
X-link: facebookresearch/FBGEMM#2119


If the input tensors use a device that differs from the current device, it would cause the wrong device to be used for things such as workspace allocation (when using `cutlass::device_memory::allocation`) and kernel to run on the wrong stream. Either would break the kernel. As a fix we add the `CUDAGuard` to ensure correct device is used.
- `cutlass::device_memory::allocation` is a wrapper around [`cudaMalloc`](https://github.com/NVIDIA/cutlass/blob/2252254ce2c3f11ef5cfff9721ebbe7bd62cf8cb/tools/util/include/cutlass/util/device_memory.h#L56), but this would bypass PyTorch CCA. We replace all usages with torch tensor allocation instead which would be less error prone and allow proper memory reuse.

Differential Revision: D86768064

Author: cthi

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped.cu

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 
 #include "bf16bf16bf16_grouped/bf16bf16bf16_grouped_manifest.cuh"
 #include "fbgemm_gpu/quantize/tuning_cache.cuh"
@@ -758,6 +759,8 @@ at::Tensor dispatch_bf16_grouped_kernel(
 
 template <typename OutputType>
 OutputType _bf16bf16bf16_grouped(at::TensorList X, at::TensorList W) {
+  c10::cuda::CUDAGuard deviceGuard(X[0].device());
+
   at::Tensor Y;
   int64_t total_M = 0;
   int64_t G = X.size();
@@ -816,6 +819,8 @@ at::Tensor bf16bf16bf16_grouped_stacked(
     at::Tensor M_sizes,
     std::optional<at::Tensor> out,
     std::optional<int64_t> num_sms) {
+  c10::cuda::CUDAGuard deviceGuard(X.device());
+
   int64_t total_M = X.size(0);
   int64_t N = W.size(1);
   int64_t K = W.size(2);
@@ -853,6 +858,8 @@ at::Tensor bf16bf16bf16_grouped_dynamic(
   TORCH_CHECK(
       zero_start_index_M.device() == X.device(),
       "zero_start_index_M must be on same device as inputs.");
+  c10::cuda::CUDAGuard deviceGuard(X.device());
+
   int64_t G = X.size(0);
   int64_t M = X.size(1);
   int64_t N = W.size(1);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_grad/bf16bf16bf16_grouped_grad_common.cuh

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <cutlass/util/device_memory.h>
 #include <cutlass/util/packed_stride.hpp>
 
@@ -112,6 +113,8 @@ at::Tensor bf16bf16bf16_grouped_grad_impl(
     at::Tensor output,
     int sm_count,
     std::optional<at::Tensor> M_sizes) {
+  c10::cuda::CUDAGuard deviceGuard(X.device());
+
   int64_t G;
   at::TensorOptions options;
   G = W.size(0);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_common.cuh

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <cutlass/util/device_memory.h>
 #include <cutlass/util/packed_stride.hpp>
 
@@ -113,6 +114,8 @@ at::Tensor bf16bf16bf16_grouped_wgrad_impl(
     at::Tensor M_sizes,
     at::Tensor output,
     int sm_count) {
+  c10::cuda::CUDAGuard deviceGuard(X.device());
+
   int64_t G;
   at::TensorOptions options;
   G = M_sizes.size(0);
@@ -377,6 +380,8 @@ at::Tensor bf16bf16bf16_grouped_wgrad_sm100_impl(
     at::Tensor M_sizes,
     at::Tensor output,
     int sm_count) {
+  c10::cuda::CUDAGuard deviceGuard(X.device());
+
   int64_t G;
   at::TensorOptions options;
   G = M_sizes.size(0);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16i4bf16.cu

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 
 #include "cutlass/cutlass.h"
 
@@ -41,6 +42,8 @@ at::Tensor _bf16i4bf16(
     at::Tensor w_scale_group,
     at::Tensor w_zero_group,
     at::Tensor Y) {
+  c10::cuda::CUDAGuard deviceGuard(X.device());
+
   // Get shape information from input tensors.
   int M = size_to_dim_(X.dim() - 1, X.sizes());
   int K = X.size(-1);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16i4bf16_rowwise_batched.cu

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <cutlass/util/device_memory.h>
 #include <cutlass/util/packed_stride.hpp>
 
@@ -48,6 +49,8 @@ at::Tensor bf16i4bf16_rowwise_batched_impl(
     at::Tensor WQ, // INT4
     at::Tensor w_scale,
     at::Tensor w_zp) {
+  c10::cuda::CUDAGuard deviceGuard(X.device());
+
   // XQ: B x M x K
   // WQ: B x N x K
   // output: B x M x N
@@ -244,7 +247,8 @@ at::Tensor bf16i4bf16_rowwise_batched_impl(
   size_t workspace_size = Gemm::get_workspace_size(arguments);
 
   // Allocate workspace memory
-  cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
+  at::Tensor workspace =
+      at::empty(workspace_size, X.options().dtype(at::kByte));
 
   // Check the problem size is supported or not
   cutlass::Status status = gemm.can_implement(arguments);
@@ -253,7 +257,7 @@ at::Tensor bf16i4bf16_rowwise_batched_impl(
   }
 
   // Initialize CUTLASS kernel with arguments and workspace pointer
-  status = gemm.initialize(arguments, workspace.get());
+  status = gemm.initialize(arguments, workspace.data_ptr());
   if (status != cutlass::Status::kSuccess) {
     throw std::runtime_error("cutlass cannot initialize");
   }

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16i4bf16_shuffled_grouped.cu

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 
 #include "cutlass/cutlass.h"
 
@@ -137,6 +138,8 @@ void _bf16i4bf16_shuffled_grouped(
     at::Tensor w_zero_group,
     at::Tensor M_sizes,
     at::Tensor Y) {
+  c10::cuda::CUDAGuard deviceGuard(X.device());
+
   // Get basic shape information.
   int G = M_sizes.size(0);
   // X is shape [total_M, K]

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16/f4f4bf16_common.cuh

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <cutlass/util/device_memory.h>
 #include <cutlass/util/packed_stride.hpp>
 
@@ -34,6 +35,8 @@ at::Tensor _f4f4bf16(
     at::Tensor x_scale,
     at::Tensor w_scale,
     std::optional<at::Tensor> global_scale) {
+  c10::cuda::CUDAGuard deviceGuard(XQ.device());
+
   int M = XQ.size(0);
   int N = WQ.size(0);
   int K = XQ.size(1) * 2; // Since K is packed
@@ -214,7 +217,8 @@ at::Tensor _f4f4bf16(
   size_t workspace_size = Gemm::get_workspace_size(arguments);
 
   // Allocate workspace memory
-  cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
+  at::Tensor workspace =
+      at::empty(workspace_size, XQ.options().dtype(at::kByte));
 
   // Check the problem size is supported or not
   cutlass::Status status = gemm.can_implement(arguments);
@@ -223,7 +227,7 @@ at::Tensor _f4f4bf16(
   }
 
   // Initialize CUTLASS kernel with arguments and workspace pointer
-  status = gemm.initialize(arguments, workspace.get());
+  status = gemm.initialize(arguments, workspace.data_ptr());
   if (status != cutlass::Status::kSuccess) {
     throw std::runtime_error("cutlass cannot initialize");
   }

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16_grouped/f4f4bf16_grouped_common.cuh

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <cutlass/util/device_memory.h>
 #include <cutlass/util/packed_stride.hpp>
 
@@ -160,6 +161,8 @@ at::Tensor f4f4bf16_grouped_impl(
     std::optional<at::Tensor> M_sizes,
     std::optional<at::Tensor> global_scale,
     std::optional<at::Tensor> starting_row_after_padding) {
+  c10::cuda::CUDAGuard deviceGuard(XQ.device());
+
   // The number of groups the kernel uses may vary.
   const int64_t G = [&]() {
     if (M_sizes) {

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16.cu

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <cutlass/util/host_tensor.h>
 #include <cutlass/util/packed_stride.hpp>
 
@@ -36,6 +37,7 @@ at::Tensor f8f8bf16_impl(
     at::Tensor XQ, // FP8
     at::Tensor WQ, // FP8
     at::Tensor scale) {
+  c10::cuda::CUDAGuard deviceGuard(XQ.device());
   // XQ: M x K
   // WQ: N x K
   // output: M x N
@@ -179,7 +181,8 @@ at::Tensor f8f8bf16_impl(
   size_t workspace_size = Gemm::get_workspace_size(arguments);
 
   // Allocate workspace memory
-  cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
+  at::Tensor workspace =
+      at::empty(workspace_size, XQ.options().dtype(at::kByte));
 
   // Check the problem size is supported or not
   cutlass::Status status = gemm.can_implement(arguments);
@@ -188,7 +191,7 @@ at::Tensor f8f8bf16_impl(
   }
 
   // Initialize CUTLASS kernel with arguments and workspace pointer
-  status = gemm.initialize(arguments, workspace.get());
+  status = gemm.initialize(arguments, workspace.data_ptr());
   if (status != cutlass::Status::kSuccess) {
     throw std::runtime_error("cutlass cannot initialize");
   }

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_blockwise.cu

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <cutlass/util/host_tensor.h>
 #include <cutlass/util/packed_stride.hpp>
 
@@ -46,6 +47,7 @@ at::Tensor f8f8bf16_blockwise_impl(
     at::Tensor WQ, // FP8
     at::Tensor x_scale,
     at::Tensor w_scale) {
+  c10::cuda::CUDAGuard deviceGuard(XQ.device());
   // XQ: M x K
   // WQ: N x K
   // output: M x N
@@ -214,7 +216,8 @@ at::Tensor f8f8bf16_blockwise_impl(
   size_t workspace_size = Gemm::get_workspace_size(arguments);
 
   // Allocate workspace memory
-  cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
+  at::Tensor workspace =
+      at::empty(workspace_size, XQ.options().dtype(at::kByte));
 
   // Check the problem size is supported or not
   cutlass::Status status = gemm.can_implement(arguments);
@@ -223,7 +226,7 @@ at::Tensor f8f8bf16_blockwise_impl(
   }
 
   // Initialize CUTLASS kernel with arguments and workspace pointer
-  status = gemm.initialize(arguments, workspace.get());
+  status = gemm.initialize(arguments, workspace.data_ptr());
   if (status != cutlass::Status::kSuccess) {
     throw std::runtime_error("cutlass cannot initialize");
   }