Apply quant layer norm fixes from vllm-project#27865, inv scale fix for int8

Signed-off-by: ElizaWszola <ewszola@redhat.com>

Author: ElizaWszola

csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

@@ -130,15 +130,14 @@ __global__ void rms_norm_per_block_quant_kernel_3(
   // RMS Norm + Quant
   int token_idx = blockIdx.x * hidden_size / group_size;
   if constexpr (std::is_same_v<scalar_out_t, int8_t>) {
-    for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
-      auto token_group_idx = token_idx + i / group_size;
-      token_scale[token_group_idx] = 1.0f / token_scale[token_group_idx];
-    }
+    // Don't invert token_scale here: do it inside the norm_and_quant kernel
+    // We do it because particular elements of token_scale can be shared
+    // between multiple threads, so this way, we avoid extra synchronization
+    // overhead.
     vllm::norm_and_quant<scalar_t, scalar_out_t, true, has_residual>(
         out, input, weight, rms[blockIdx.x], token_scale + token_idx,
         hidden_size, residual, group_size);
   } 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[blockIdx.x], token_scale + token_idx,
         hidden_size, residual, group_size);

csrc/quantization/fused_kernels/layernorm_utils.cuh

@@ -148,7 +148,11 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
     // Norm
     x = static_cast<float>(static_cast<scalar_t>(x * rms) * weight[i]);
     // Quant
-    auto scale_val = (group_size > 0 ? scale[i / group_size] : *scale);
+    // If groupwise is_scale_inverted is true, so we invert the scale here.
+    auto scale_val =
+        (group_size > 0 ? (is_scale_inverted ? 1.0f / scale[i / group_size]
+                                             : scale[i / group_size])
+                        : *scale);
     output[token_offset + i] =
         ScaledQuant<scalar_out_t, is_scale_inverted>::quant_fn(x, scale_val);
   }

csrc/quantization/w8a8/int8/scaled_quant.cu

@@ -1,5 +1,6 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <torch/all.h>
+#include <c10/cuda/CUDAGuard.h>
 
 #include <cmath>
 
@@ -275,6 +276,7 @@ void static_scaled_int8_quant(torch::Tensor& out,          // [..., hidden_size]
   int const num_tokens = input.numel() / hidden_size;
   dim3 const grid(num_tokens);
   dim3 const block(std::min(hidden_size, 256));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   VLLM_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "static_scaled_int8_quant_kernel", [&] {
@@ -306,6 +308,7 @@ void dynamic_scaled_int8_quant(
   int const num_tokens = input.numel() / hidden_size;
   dim3 const grid(num_tokens);
   dim3 const block(std::min(hidden_size, 256));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   VLLM_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "dynamic_scaled_int8_quant_kernel", [&] {

tests/kernels/core/test_fused_quant_layernorm.py

@@ -17,7 +17,7 @@
 
 DTYPES = [torch.bfloat16, torch.float]
 QUANT_DTYPES = [torch.int8, torch.float8_e4m3fn]
-VEC_HIDDEN_SIZES = range(1024, 1030)
+VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
 # Avoid combinatorial explosion with full Cartesian product
 NUM_TOKENS_HIDDEN_SIZES = [
     *[(1, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5120, 5137]],
@@ -165,8 +165,8 @@ def test_rms_norm(
         # skip
         return
 
-    # blockwise baseline doesn't support scale_ub
     if group_size is not None and has_scale_ub:
+        # blockwise baseline doesn't support scale_ub
         return
 
     if has_scale_ub and quant_dtype != torch.float8_e4m3fn:
@@ -197,14 +197,31 @@ def test_rms_norm(
 
     assert ref_out.dtype == quant_dtype
     assert ops_out.dtype == quant_dtype
-    assert torch.allclose(ref_scales, ops_scales)
     if quant_dtype == torch.int8:
+        assert torch.allclose(ref_scales, ops_scales, atol=1e-6)
         # big atol to account for round-off errors.
         assert torch.allclose(ref_out, ops_out, atol=1)
     else:
-        assert torch.allclose(
-            ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32)
-        )
+        assert torch.allclose(ref_scales, ops_scales)
+        a = ref_out.to(dtype=torch.float32)
+        b = ops_out.to(dtype=torch.float32)
+        ok = torch.allclose(a, b, atol=1e-6)
+        if not ok:
+            # fallback: compare dequantized values with relaxed tolerance
+            if group_size is None:
+                a_deq = a * ref_scales.view(-1, 1)
+                b_deq = b * ops_scales.view(-1, 1)
+            else:
+                a_deq = a * ref_scales.repeat_interleave(group_size[1], dim=1)
+                b_deq = b * ops_scales.repeat_interleave(group_size[1], dim=1)
+            # NOTE: It is possible that some future test cases trigger this
+            # max diff due to precision issues. If such an error is
+            # encountered, it's recommended to inspect the differences between
+            # all corresponding elements from each tensor (e.g. by looping over
+            # them) and checking how many the max diff error shows up on (just
+            # a few bad elements should still be considered acceptable).
+            ok = torch.allclose(a_deq, b_deq, rtol=5e-2, atol=5e-2)
+        assert ok
     if add_residual:
         assert torch.allclose(ref_residual, ops_residual)