dynamic quant layernorm

Signed-off-by: charlifu <charlifu@amd.com>

Author: charlifu

csrc/ops.h

@@ -115,6 +115,7 @@ void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
                                       torch::Tensor& scales,
                                       double const epsilon,
                                       std::optional<torch::Tensor> scale_ub,
+                                      std::optional<torch::Tensor> residual_out,
                                       std::optional<torch::Tensor> residual);
 
 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,

csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

@@ -15,6 +15,7 @@ __device__ void rms_norm_dynamic_per_token_quant_vec(
     scalar_t const* __restrict__ input,   // [..., hidden_size]
     scalar_t const* __restrict__ weight,  // [hidden_size]
     float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
+    scalar_t* __restrict__ residual_out = nullptr,
     scalar_t* __restrict__ residual = nullptr) {
   float rms = 0.0f;
   float token_scale = 0.0f;
@@ -33,12 +34,14 @@ __device__ void rms_norm_dynamic_per_token_quant_vec(
   if constexpr (std::is_same_v<scalar_out_t, int8_t>) {
     vllm::vectorized::norm_and_quant<scalar_t, scalar_out_t, true,
                                      has_residual>(
-        out, input, weight, rms, 1.0f / token_scale, hidden_size, residual);
+        out, input, weight, rms, 1.0f / token_scale, hidden_size, residual_out,
+        residual);
   } else {
     // FP8 - Do not invert token_scale for exact match with FBGemm
     vllm::vectorized::norm_and_quant<scalar_t, scalar_out_t, false,
-                                     has_residual>(
-        out, input, weight, rms, token_scale, hidden_size, residual);
+                                     has_residual>(out, input, weight, rms,
+                                                   token_scale, hidden_size,
+                                                   residual_out, residual);
   }
 }
 
@@ -50,6 +53,7 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
     scalar_t const* __restrict__ input,   // [..., hidden_size]
     scalar_t const* __restrict__ weight,  // [hidden_size]
     float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
+    scalar_t* __restrict__ residual_out = nullptr,
     scalar_t* __restrict__ residual = nullptr) {
   // For vectorization, token_input and token_output pointers need to be
   // aligned at 8-byte and 4-byte addresses respectively.
@@ -76,11 +80,13 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
   // RMS Norm + Quant
   if constexpr (std::is_same_v<scalar_out_t, int8_t>) {
     vllm::norm_and_quant<scalar_t, scalar_out_t, true, has_residual>(
-        out, input, weight, rms, 1.0f / token_scale, hidden_size, residual);
+        out, input, weight, rms, 1.0f / token_scale, hidden_size, residual_out,
+        residual);
   } else {
     // FP8 - Do not invert s_token_scale for exact match with FBGemm
     vllm::norm_and_quant<scalar_t, scalar_out_t, false, has_residual>(
-        out, input, weight, rms, token_scale, hidden_size, residual);
+        out, input, weight, rms, token_scale, hidden_size, residual_out,
+        residual);
   }
 }
 }  // namespace vllm
@@ -94,6 +100,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
     torch::Tensor& scales,        // [num_tokens]
     double const var_epsilon,     // Variance epsilon used in norm calculation
     std::optional<at::Tensor> const& scale_ub,
+    std::optional<at::Tensor>& residual_out,
     std::optional<at::Tensor>& residual) {
   int32_t hidden_size = input.size(-1);
   auto num_tokens = input.numel() / hidden_size;
@@ -112,7 +119,9 @@ void rms_norm_dynamic_per_token_quant_dispatch(
                   out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
                   input.data_ptr<scalar_in_t>(), weight.data_ptr<scalar_in_t>(),
                   scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
-                  var_epsilon, hidden_size, residual->data_ptr<scalar_in_t>());
+                  var_epsilon, hidden_size,
+                  residual_out->data_ptr<scalar_in_t>(),
+                  residual->data_ptr<scalar_in_t>());
         });
 
   } else {
@@ -124,7 +133,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
                   out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
                   input.data_ptr<scalar_in_t>(), weight.data_ptr<scalar_in_t>(),
                   scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
-                  var_epsilon, hidden_size, nullptr);
+                  var_epsilon, hidden_size, nullptr, nullptr);
         });
   }
 }
@@ -135,13 +144,18 @@ void rms_norm_dynamic_per_token_quant(
     torch::Tensor const& weight,  // [hidden_size]
     torch::Tensor& scales,        // [num_tokens]
     double const var_epsilon,     // Variance epsilon used in norm calculation
-    std::optional<at::Tensor> scale_ub, std::optional<at::Tensor> residual) {
+    std::optional<at::Tensor> scale_ub, std::optional<at::Tensor> residual_out,
+    std::optional<at::Tensor> residual) {
   static c10::ScalarType kFp8Type = is_fp8_ocp()
                                         ? c10::ScalarType::Float8_e4m3fn
                                         : c10::ScalarType::Float8_e4m3fnuz;
   TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
   TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
 
+  if (residual.has_value()) {
+    TORCH_CHECK(residual_out.has_value());
+  }
+
   if (scale_ub.has_value()) {
     TORCH_CHECK(out.dtype() == kFp8Type);
   }
@@ -150,6 +164,7 @@ void rms_norm_dynamic_per_token_quant(
   VLLM_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "rms_norm_dynamic_per_token_quant_dispatch", [&] {
         rms_norm_dynamic_per_token_quant_dispatch<scalar_t>(
-            out, input, weight, scales, var_epsilon, scale_ub, residual);
+            out, input, weight, scales, var_epsilon, scale_ub, residual_out,
+            residual);
       });
 }

csrc/quantization/fused_kernels/layernorm_utils.cuh

@@ -100,6 +100,7 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
                                scalar_t const* __restrict__ weight,
                                float const rms, float const scale,
                                int32_t const hidden_size,
+                               scalar_t* __restrict__ residual_out = nullptr,
                                scalar_t* __restrict__ residual = nullptr) {
   int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
   ;
@@ -108,7 +109,7 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
     float x = static_cast<float>(input[token_offset + i]);
     if constexpr (has_residual) {
       x += static_cast<float>(residual[token_offset + i]);
-      residual[token_offset + i] = static_cast<scalar_t>(x);
+      residual_out[token_offset + i] = static_cast<scalar_t>(x);
     }
     // Norm
     x = static_cast<float>(static_cast<scalar_t>(x * rms) * weight[i]);
@@ -268,6 +269,7 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
                                scalar_t const* __restrict__ weight,
                                float const rms, float const scale,
                                int32_t const hidden_size,
+                               scalar_t* __restrict__ residual_out = nullptr,
                                scalar_t* __restrict__ residual = nullptr) {
   int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
   ;
@@ -279,8 +281,11 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
       reinterpret_cast<vec4_t<scalar_t> const*>(weight);
   q8x4_t<scalar_out_t>* vec_output =
       reinterpret_cast<q8x4_t<scalar_out_t>*>(&output[token_offset]);
+  vec4_t<scalar_t>* vec_residual_out = nullptr;
   vec4_t<scalar_t>* vec_residual = nullptr;
   if constexpr (has_residual) {
+    vec_residual_out =
+        reinterpret_cast<vec4_t<scalar_t>*>(&residual_out[token_offset]);
     vec_residual = reinterpret_cast<vec4_t<scalar_t>*>(&residual[token_offset]);
   }
 
@@ -311,7 +316,7 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
       for (int j = 0; j < VEC_SIZE; ++j) {
         r.val[j] = static_cast<scalar_t>(x.val[j]);
       }
-      vec_residual[i] = r;
+      vec_residual_out[i] = r;
     }
 
     q8x4_t<scalar_out_t> out;

csrc/torch_bindings.cpp

@@ -199,7 +199,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def(
       "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, "
       "Tensor weight, Tensor! scale, float epsilon, "
-      "Tensor? scale_ub, Tensor!? residual) -> ()");
+      "Tensor? scale_ub, Tensor!? residual_out, Tensor? residual) -> ()");
   ops.impl("rms_norm_dynamic_per_token_quant", torch::kCUDA,
            &rms_norm_dynamic_per_token_quant);
 

vllm/compilation/fusion.py

@@ -400,6 +400,7 @@ def replacement(result: torch.Tensor, result_rms: torch.Tensor,
                                      scale=scale,
                                      epsilon=self.epsilon,
                                      scale_ub=None,
+                                     residual_out=None,
                                      residual=None)
 
             # result, scale
@@ -471,11 +472,14 @@ def __init__(self,
     def register(self, pm_pass: PatternMatcherPass,
                  record_match: Callable[[MultiOutputMatch], bool]):
 
-        def pattern(result: torch.Tensor, input: torch.Tensor,
+        def pattern(result: torch.Tensor, result_rms: torch.Tensor,
+                    input: torch.Tensor, residual_out: torch.Tensor,
                     residual: torch.Tensor, weight: torch.Tensor,
                     scale: torch.Tensor):
             at = auto_functionalized(RMS_ADD_OP,
+                                     result=result_rms,
                                      input=input,
+                                     residual_out=residual_out,
                                      residual=residual,
                                      weight=weight,
                                      epsilon=self.epsilon)
@@ -488,7 +492,8 @@ def pattern(result: torch.Tensor, input: torch.Tensor,
             # result, residual, scale
             return at1[1], at[2], at1[2]
 
-        def replacement(result: torch.Tensor, input: torch.Tensor,
+        def replacement(result: torch.Tensor, result_rms: torch.Tensor,
+                        input: torch.Tensor, residual_out: torch.Tensor,
                         residual: torch.Tensor, weight: torch.Tensor,
                         scale: torch.Tensor):
             at = auto_functionalized(self.FUSED_OP,
@@ -498,14 +503,17 @@ def replacement(result: torch.Tensor, input: torch.Tensor,
                                      scale=scale,
                                      epsilon=self.epsilon,
                                      scale_ub=None,
+                                     residual_out=residual_out,
                                      residual=residual)
 
             # result, residual, scale
             return at[1], at[3], at[2]
 
         inputs = [
             torch.empty(5, 4, device="cuda", dtype=self.quant_dtype),  # result
+            empty_bf16(5, 4),  # result_rms 
             empty_bf16(5, 4),  # input
+            empty_bf16(5, 4),  # residual_out
             empty_bf16(5, 4),  # residual
             empty_bf16(1, 5),  # weight
             empty_fp32(1, 1)  # scale

vllm/utils.py

@@ -2415,31 +2415,13 @@ def memory_profiling(
     result.weights_memory = weights_memory
 
     result.before_profile.measure()
-    logger.info("Before Profiling:", "torch_peak:",
-                result.before_profile.torch_peak / GiB_bytes, "torch_mem:",
-                result.before_profile.torch_memory / GiB_bytes,
-                "non_torch_mem:",
-                result.before_profile.non_torch_memory / GiB_bytes,
-                "cuda_mem:", result.before_profile.cuda_memory / GiB_bytes,
-                "free_mem:", result.before_profile.free_memory / GiB_bytes,
-                "total_mem:", result.before_profile.total_memory / GiB_bytes)
-    torch.cuda.memory._record_memory_history()
 
     yield result
 
     gc.collect()
     torch.cuda.empty_cache()
 
     result.after_profile.measure()
-    logger.info("After Profiling: torch_peak:",
-                result.after_profile.torch_peak / GiB_bytes, "torch_mem:",
-                result.after_profile.torch_memory / GiB_bytes,
-                "non_torch_mem:",
-                result.after_profile.non_torch_memory / GiB_bytes, "cuda_mem:",
-                result.after_profile.cuda_memory / GiB_bytes, "free_mem:",
-                result.after_profile.free_memory / GiB_bytes, "total_mem:",
-                result.after_profile.total_memory / GiB_bytes)
-
     diff_profile = result.after_profile - result.before_profile
     diff_from_create = result.after_profile - result.before_create
     result.torch_peak_increase = diff_profile.torch_peak