support rms_norm_dynamic_per_token_quant

CMakeLists.txt

@@ -261,7 +261,7 @@ set(VLLM_EXT_SRC
   # "csrc/quantization/gptq/q_gemm.cu"
   "csrc/quantization/compressed_tensors/int8_quant_kernels.cu"
   # "csrc/quantization/fp8/common.cu"
-  # "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
+  "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
   "csrc/quantization/gguf/gguf_kernel.cu"
   # "csrc/quantization/activation_kernels.cu"
   "csrc/cuda_utils_kernels.cu"

csrc/ops.h

@@ -233,13 +233,13 @@ void apply_repetition_penalties_(torch::Tensor& logits,
 //                                          torch::Tensor& weight,
 //                                          torch::Tensor& scale, double epsilon);
 
-// void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
-//                                       torch::Tensor const& input,
-//                                       torch::Tensor const& weight,
-//                                       torch::Tensor& scales,
-//                                       double const epsilon,
-//                                       std::optional<torch::Tensor> scale_ub,
-//                                       std::optional<torch::Tensor> residual);
+void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
+                                      torch::Tensor const& input,
+                                      torch::Tensor const& weight,
+                                      torch::Tensor& scales,
+                                      double const epsilon,
+                                      std::optional<torch::Tensor> scale_ub,
+                                      std::optional<torch::Tensor> residual);
 
 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                       std::optional<torch::Tensor> key, int64_t head_size,

csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

@@ -6,6 +6,20 @@
 #include "layernorm_utils.cuh"
 #include "quant_conversions.cuh"
 
+// Determines the preferred FP8 type for the current platform.
+// Note that for CUDA this just returns true,
+// but on ROCm it will check device props.
+static bool is_fp8_ocp() {
+#ifndef USE_ROCM
+  return true;
+#else
+  auto dprops = at::cuda::getCurrentDeviceProperties();
+  std::string device_arch = dprops->gcnArchName;
+  size_t substring = device_arch.find("gfx94");
+  return substring == std::string::npos;
+#endif
+}
+
 namespace vllm {
 
 template <typename scalar_t, typename scalar_out_t, bool has_residual = false>

csrc/torch_bindings.cpp

@@ -372,12 +372,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 //            &fused_add_rms_norm_static_fp8_quant);
 
   // Fused Layernorm + Quant kernels
-//   ops.def(
-//       "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, "
-//       "Tensor weight, Tensor! scale, float epsilon, "
-//       "Tensor? scale_ub, Tensor!? residual) -> ()");
-//   ops.impl("rms_norm_dynamic_per_token_quant", torch::kCUDA,
-//            &rms_norm_dynamic_per_token_quant);
+  ops.def(
+      "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, "
+      "Tensor weight, Tensor! scale, float epsilon, "
+      "Tensor? scale_ub, Tensor!? residual) -> ()");
+  ops.impl("rms_norm_dynamic_per_token_quant", torch::kCUDA,
+           &rms_norm_dynamic_per_token_quant);
 
   // Rotary embedding
   // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.