Sync updates for CUDA 13 compat

build.toml

@@ -95,6 +95,7 @@ depends = ["torch"]
 include = ["."]
 src = [
     "compressed_tensors/int8_quant_kernels.cu",
+    "cub_helpers.h",
     "dispatch_utils.h",
     "vectorization_utils.cuh",
 ]
@@ -119,6 +120,7 @@ include = ["."]
 src = [
     "fp8/common.cu",
     "fp8/common.cuh",
+    "cub_helpers.h",
     "dispatch_utils.h",
     "utils.cuh",
     "vectorization.cuh",

compressed_tensors/int8_quant_kernels.cu

@@ -3,6 +3,7 @@
 
 #include <cmath>
 
+#include "../cub_helpers.h"
 #include "../dispatch_utils.h"
 #include "../vectorization_utils.cuh"
 
@@ -168,7 +169,7 @@ __global__ void dynamic_scaled_int8_quant_kernel(
   }
   using BlockReduce = cub::BlockReduce<float, 256>;
   __shared__ typename BlockReduce::TempStorage tmp;
-  float block_max = BlockReduce(tmp).Reduce(thread_max, cub::Max{}, blockDim.x);
+  float block_max = BlockReduce(tmp).Reduce(thread_max, CubMaxOp{}, blockDim.x);
   __shared__ float absmax;
   if (tid == 0) {
     absmax = block_max;

cub_helpers.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#ifndef USE_ROCM
+  #include <cub/cub.cuh>
+  #if CUB_VERSION >= 200800
+    #include <cuda/std/functional>
+using CubAddOp = cuda::std::plus<>;
+using CubMaxOp = cuda::maximum<>;
+  #else   // if CUB_VERSION < 200800
+using CubAddOp = cub::Sum;
+using CubMaxOp = cub::Max;
+  #endif  // CUB_VERSION
+#else
+  #include <hipcub/hipcub.hpp>
+using CubAddOp = cub::Sum;
+using CubMaxOp = cub::Max;
+#endif  // USE_ROCM

fp8/common.cu

@@ -1,5 +1,6 @@
 #include "common.cuh"
 #include "dispatch_utils.h"
+#include "../cub_helpers.h"
 
 #include <c10/cuda/CUDAGuard.h>
 
@@ -55,7 +56,7 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel(
   using BlockReduce = cub::BlockReduce<float, 256>;
   __shared__ typename BlockReduce::TempStorage reduceStorage;
   float const block_absmax_val_maybe =
-      BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim.x);
+      BlockReduce(reduceStorage).Reduce(absmax_val, CubMaxOp{}, blockDim.x);
   __shared__ float token_scale;
   if (tid == 0) {
     if (scale_ub) {