diff --git a/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py b/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py
index c84fd5dda84e..c4daeb34373a 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py
@@ -254,6 +254,8 @@ def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl(
     Returns:
         Output tensor of shape (batch_size, output_dim) in bfloat16 format
     """
+    if input.is_cuda:
+        torch.cuda.set_device(input.device)
 
     def run_flashinfer_deepgemm_swapAB(
         input: torch.Tensor,