Use PDL

Signed-off-by: yewentao256 <zhyanwentao@126.com>

csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu

@@ -6,6 +6,7 @@
 #include "libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h"
 
 #include <cmath>
+#include <cstdint>
 
 #ifdef USE_ROCM
   #include <hip/hip_fp8.h>
@@ -18,6 +19,14 @@
 #include "libtorch_stable/dispatch_utils.h"
 #include "libtorch_stable/torch_utils.h"
 
+namespace {
+
+inline bool Is16ByteAligned(const void* ptr) {
+  return (reinterpret_cast<std::uintptr_t>(ptr) & 0xFu) == 0;
+}
+
+}  // namespace
+
 __device__ __forceinline__ float GroupReduceMax(float val) {
 #ifdef USE_ROCM
   // 16-thread logical groups may pack up to four per 64-lane wavefront; use a
@@ -95,6 +104,80 @@ __device__ __forceinline__ void QuantizeGroup(
       scalar_op_quant);   // scalar handler
 }
 
+template <typename T, int VEC_SIZE>
+__device__ __forceinline__ float LoadRegisterGroupAndComputeAbsmax(
+    const T* __restrict__ group_input, const bool is_valid_group,
+    T (&regs)[VEC_SIZE], const float eps) {
+  static_assert(VEC_SIZE * sizeof(T) == 32,
+                "LoadRegisterGroupAndComputeAbsmax assumes two uint4 loads.");
+  float local_absmax = eps;
+  if (!is_valid_group) {
+    return local_absmax;
+  }
+
+  uint4* dst = reinterpret_cast<uint4*>(&regs[0]);
+  const uint4* src = reinterpret_cast<const uint4*>(group_input);
+  dst[0] = src[0];
+  dst[1] = src[1];
+#pragma unroll
+  for (int i = 0; i < VEC_SIZE; ++i) {
+    float v = fabsf(static_cast<float>(regs[i]));
+    local_absmax = fmaxf(local_absmax, v);
+  }
+
+  return local_absmax;
+}
+
+__device__ __forceinline__ float GroupReduceMax8(float val) {
+#ifdef USE_ROCM
+  const int lane_in_wave = threadIdx.x % warpSize;
+  const unsigned long long mask = 0xFFull << (lane_in_wave & ~7);
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 4, 8));
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 2, 8));
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 1, 8));
+#else
+  unsigned mask = 0xffu << (threadIdx.x & 24u);
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 4));
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 2));
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 1));
+#endif
+  return val;
+}
+
+template <typename T, typename DST_DTYPE, int VEC_SIZE, bool SCALE_IS_INVERTED>
+__device__ __forceinline__ uint4
+QuantizeRegisterGroup(const T (&regs)[VEC_SIZE], const float scale,
+                      const float min_8bit, const float max_8bit) {
+  uint32_t packed_lo = 0;
+  uint32_t packed_lo_hi = 0;
+  uint32_t packed_hi_lo = 0;
+  uint32_t packed_hi = 0;
+#pragma unroll
+  for (int i = 0; i < VEC_SIZE; ++i) {
+    float q = static_cast<float>(regs[i]);
+    if constexpr (SCALE_IS_INVERTED) {
+      q *= scale;
+    } else {
+      q /= scale;
+    }
+    q = fminf(fmaxf(q, min_8bit), max_8bit);
+    DST_DTYPE qb = DST_DTYPE(q);
+    uint8_t byte = *reinterpret_cast<uint8_t*>(&qb);
+    const int shift = (i & 3) * 8;
+    if (i < 4) {
+      packed_lo |= static_cast<uint32_t>(byte) << shift;
+    } else if (i < 8) {
+      packed_lo_hi |= static_cast<uint32_t>(byte) << shift;
+    } else if (i < 12) {
+      packed_hi_lo |= static_cast<uint32_t>(byte) << shift;
+    } else {
+      packed_hi |= static_cast<uint32_t>(byte) << shift;
+    }
+  }
+
+  return make_uint4(packed_lo, packed_lo_hi, packed_hi_lo, packed_hi);
+}
+
 template <typename T, typename DST_DTYPE, bool IS_COLUMN_MAJOR = false,
           bool SCALE_UE8M0 = false, typename scale_packed_t = float>
 __global__ void per_token_group_quant_8bit_kernel(
@@ -190,6 +273,96 @@ inline int GetGroupsPerBlockX(int64_t padded_groups_per_row) {
   return 4;
 }
 
+template <typename T, typename DST_DTYPE, bool IS_COLUMN_MAJOR,
+          bool SCALE_UE8M0, int GROUP_SIZE, int kGroupsPerBlockX,
+          int kRowsPerBlock>
+__global__ void per_token_group_quant_8bit_register_kernel(
+    const T* __restrict__ input, void* __restrict__ output_q,
+    float* __restrict__ output_s, const int groups_per_row, const int mn,
+    const float eps, const float min_8bit, const float max_8bit,
+    const int scale_stride) {
+  static_assert(GROUP_SIZE == 128, "fast path supports GROUP_SIZE==128 only");
+  constexpr int THREADS_PER_GROUP = 8;
+  constexpr int VEC_SIZE = 32 / sizeof(T);  // 16 for bf16/fp16
+  static_assert(GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE,
+                "GROUP_SIZE must equal THREADS_PER_GROUP * VEC_SIZE");
+  static_assert(32 % THREADS_PER_GROUP == 0,
+                "THREADS_PER_GROUP must divide warp size for the shuffle "
+                "mask to be valid");
+  static_assert(
+      kGroupsPerBlockX > 0 && (kGroupsPerBlockX & (kGroupsPerBlockX - 1)) == 0,
+      "kGroupsPerBlockX must be a positive power of 2");
+  static_assert(kRowsPerBlock > 0, "kRowsPerBlock must be positive");
+
+  const int local_group_id = threadIdx.x / THREADS_PER_GROUP;
+  const int lane_id = threadIdx.x % THREADS_PER_GROUP;
+
+  const int sf_k_local = local_group_id % kGroupsPerBlockX;
+  const int row_local = local_group_id / kGroupsPerBlockX;
+  const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local;
+  const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  cudaGridDependencySynchronize();
+#endif
+
+  if (mn_idx >= mn) {
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
+    return;
+  }
+
+  const bool is_valid_group = sf_k_idx < groups_per_row;
+
+  // register-only fast path for group_size == 128.
+  alignas(16) T regs[VEC_SIZE];
+  const T* group_input =
+      input + static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
+      sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
+  // each 8-lane subgroup processes one logical group.
+  float local_absmax = LoadRegisterGroupAndComputeAbsmax<T, VEC_SIZE>(
+      group_input, is_valid_group, regs, eps);
+  local_absmax = GroupReduceMax8(local_absmax);
+
+  float y_s = local_absmax / max_8bit;
+  if constexpr (SCALE_UE8M0) {
+    y_s = exp2f(ceilf(log2f(fmaxf(fabsf(y_s), 1e-10f))));
+  }
+
+  if (lane_id == 0 && is_valid_group) {
+    float* scale_output;
+    if constexpr (IS_COLUMN_MAJOR) {
+      scale_output =
+          output_s + static_cast<int64_t>(sf_k_idx) * scale_stride + mn_idx;
+    } else {
+      scale_output =
+          output_s + static_cast<int64_t>(mn_idx) * groups_per_row + sf_k_idx;
+    }
+    *scale_output = y_s;
+  }
+
+  if (!is_valid_group) {
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
+    return;
+  }
+
+  // pack into 16 fp8/int8 bytes (= uint4)
+  uint4 packed_out = QuantizeRegisterGroup<T, DST_DTYPE, VEC_SIZE, false>(
+      regs, y_s, min_8bit, max_8bit);
+  DST_DTYPE* group_output =
+      static_cast<DST_DTYPE*>(output_q) +
+      static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
+      sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
+  *reinterpret_cast<uint4*>(group_output) = packed_out;
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  cudaTriggerProgrammaticLaunchCompletion();
+#endif
+}
+
 void per_token_group_quant_8bit(const torch::stable::Tensor& input,
                                 torch::stable::Tensor& output_q,
                                 torch::stable::Tensor& output_s,
@@ -218,6 +391,12 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
   const bool is_column_major = output_s.stride(0) < output_s.stride(1);
   const int scale_num_rows = output_s.size(1);
   const int scale_stride = output_s.stride(1);
+  const bool register_fast_path_is_aligned =
+      Is16ByteAligned(input.data_ptr()) && Is16ByteAligned(output_q.data_ptr());
+  const bool use_register_fast_path =
+      register_fast_path_is_aligned && group_size == 128 &&
+      (input.scalar_type() == torch::headeronly::ScalarType::Half ||
+       input.scalar_type() == torch::headeronly::ScalarType::BFloat16);
 
 #ifndef USE_ROCM
   #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES)     \
@@ -271,17 +450,121 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
     }                                                                   \
   } while (0)
 
-  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
-      input.scalar_type(), "per_token_group_quant_8bit", ([&] {
-        if (dst_type == torch::headeronly::ScalarType::Char) {
-          LAUNCH_KERNEL(scalar_t, int8_t);
-        } else {
-          VLLM_STABLE_DISPATCH_FP8_TYPES(
-              dst_type, "per_token_group_quant_8bit_fp8",
-              ([&] { LAUNCH_KERNEL(scalar_t, fp8_t); }));
-        }
-      }));
+#ifndef USE_ROCM
+  #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, IS_COLUMN_MAJOR, SCALE_UE8M0, \
+                                 KX, RY)                                     \
+    do {                                                                     \
+      cudaLaunchConfig_t config = {};                                        \
+      config.gridDim = dim3(static_cast<unsigned int>(blocks_x),             \
+                            static_cast<unsigned int>(blocks_y));            \
+      config.blockDim = dim3(num_threads_fast);                              \
+      config.dynamicSmemBytes = 0;                                           \
+      config.stream = stream;                                                \
+      cudaLaunchAttribute attrs[1];                                          \
+      attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;      \
+      attrs[0].val.programmaticStreamSerializationAllowed = 1;               \
+      config.numAttrs = 1;                                                   \
+      config.attrs = attrs;                                                  \
+      cudaLaunchKernelEx(                                                    \
+          &config,                                                           \
+          per_token_group_quant_8bit_register_kernel<                        \
+              T, DST_DTYPE, IS_COLUMN_MAJOR, SCALE_UE8M0, 128, KX, RY>,      \
+          static_cast<const T*>(input.data_ptr()), output_q.data_ptr(),      \
+          static_cast<float*>(output_s.data_ptr()),                          \
+          static_cast<int>(groups_per_row), static_cast<int>(mn),            \
+          static_cast<float>(eps), static_cast<float>(min_8bit),             \
+          static_cast<float>(max_8bit), scale_stride);                       \
+    } while (0)
+#else
+  #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, IS_COLUMN_MAJOR, SCALE_UE8M0, \
+                                 KX, RY)                                     \
+    do {                                                                     \
+      dim3 grid(static_cast<unsigned int>(blocks_x),                         \
+                static_cast<unsigned int>(blocks_y));                        \
+      dim3 block(num_threads_fast);                                          \
+      per_token_group_quant_8bit_register_kernel<                            \
+          T, DST_DTYPE, IS_COLUMN_MAJOR, SCALE_UE8M0, 128, KX, RY>           \
+          <<<grid, block, 0, stream>>>(                                      \
+              static_cast<const T*>(input.data_ptr()), output_q.data_ptr(),  \
+              static_cast<float*>(output_s.data_ptr()),                      \
+              static_cast<int>(groups_per_row), static_cast<int>(mn),        \
+              static_cast<float>(eps), static_cast<float>(min_8bit),         \
+              static_cast<float>(max_8bit), scale_stride);                   \
+    } while (0)
+#endif
 
+#define LAUNCH_REG_KERNEL(T, DST_DTYPE, IS_COLUMN_MAJOR, SCALE_UE8M0)        \
+  do {                                                                       \
+    if (kx == 16) {                                                          \
+      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, IS_COLUMN_MAJOR, SCALE_UE8M0, 16, \
+                             1);                                             \
+    } else if (kx == 8) {                                                    \
+      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, IS_COLUMN_MAJOR, SCALE_UE8M0, 8,  \
+                             2);                                             \
+    } else if (kx == 4) {                                                    \
+      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, IS_COLUMN_MAJOR, SCALE_UE8M0, 4,  \
+                             4);                                             \
+    } else {                                                                 \
+      STD_TORCH_CHECK(false, "Unsupported kx value ", kx);                   \
+    }                                                                        \
+  } while (0)
+
+#define LAUNCH_REG_KERNEL_FLAGS(T, DST_DTYPE)         \
+  do {                                                \
+    if (is_column_major) {                            \
+      if (scale_ue8m0) {                              \
+        LAUNCH_REG_KERNEL(T, DST_DTYPE, true, true);  \
+      } else {                                        \
+        LAUNCH_REG_KERNEL(T, DST_DTYPE, true, false); \
+      }                                               \
+    } else if (scale_ue8m0) {                         \
+      LAUNCH_REG_KERNEL(T, DST_DTYPE, false, true);   \
+    } else {                                          \
+      LAUNCH_REG_KERNEL(T, DST_DTYPE, false, false);  \
+    }                                                 \
+  } while (0)
+
+  if (use_register_fast_path) {
+    const int64_t k = input.size(-1);
+    const int64_t mn = input.numel() / k;
+    const int64_t groups_per_row = k / group_size;
+    const int64_t padded_groups_per_row = ((groups_per_row + 3) / 4) * 4;
+    const int kx = GetGroupsPerBlockX(padded_groups_per_row);
+    const int ry = 16 / kx;
+    const int64_t blocks_x = (mn + ry - 1) / ry;
+    const int64_t blocks_y = padded_groups_per_row / kx;
+    const int num_threads_fast = (kx * ry) * 8;
+    STD_TORCH_CHECK(
+        blocks_x <= static_cast<int64_t>(INT32_MAX) && blocks_y <= 65535,
+        "per_token_group_quant_8bit fast-path grid too large: (", blocks_x,
+        ", ", blocks_y, ").");
+
+    VLLM_STABLE_DISPATCH_HALF_TYPES(
+        input.scalar_type(), "per_token_group_quant_8bit_register", ([&] {
+          if (dst_type == torch::headeronly::ScalarType::Char) {
+            LAUNCH_REG_KERNEL_FLAGS(scalar_t, int8_t);
+          } else {
+            VLLM_STABLE_DISPATCH_FP8_TYPES(
+                dst_type, "per_token_group_quant_8bit_register_fp8",
+                ([&] { LAUNCH_REG_KERNEL_FLAGS(scalar_t, fp8_t); }));
+          }
+        }));
+  } else {
+    VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+        input.scalar_type(), "per_token_group_quant_8bit", ([&] {
+          if (dst_type == torch::headeronly::ScalarType::Char) {
+            LAUNCH_KERNEL(scalar_t, int8_t);
+          } else {
+            VLLM_STABLE_DISPATCH_FP8_TYPES(
+                dst_type, "per_token_group_quant_8bit_fp8",
+                ([&] { LAUNCH_KERNEL(scalar_t, fp8_t); }));
+          }
+        }));
+  }
+
+#undef LAUNCH_REG_KERNEL
+#undef LAUNCH_REG_KERNEL_INST
+#undef LAUNCH_REG_KERNEL_FLAGS
 #undef LAUNCH_KERNEL
 #undef LAUNCH_KERNEL_INST
 }
@@ -347,36 +630,15 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
   // alignas(16) is required so the uint4* reinterpret_cast below is
   // well-defined for T == bf16/fp16 (default alignof is 2).
   alignas(16) T regs[VEC_SIZE];
-  float local_absmax = eps;
-  if (is_valid_group) {
-    const T* group_input =
-        input + static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
-        sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
-    uint4* dst = reinterpret_cast<uint4*>(&regs[0]);
-    const uint4* src = reinterpret_cast<const uint4*>(group_input);
-    dst[0] = src[0];
-    dst[1] = src[1];
-#pragma unroll
-    for (int i = 0; i < VEC_SIZE; ++i) {
-      float v = fabsf(static_cast<float>(regs[i]));
-      local_absmax = fmaxf(local_absmax, v);
-    }
-  }
+  const T* group_input =
+      input + static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
+      sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
+  float local_absmax = LoadRegisterGroupAndComputeAbsmax<T, VEC_SIZE>(
+      group_input, is_valid_group, regs, eps);
 
   // 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the
   // 8 lanes within the warp that share a group.
-#ifdef USE_ROCM
-  const int lane_in_wave = threadIdx.x % warpSize;
-  const unsigned long long mask = 0xFFull << (lane_in_wave & ~7);
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4, 8));
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2, 8));
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1, 8));
-#else
-  unsigned mask = 0xffu << (threadIdx.x & 24u);
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4));
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2));
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1));
-#endif
+  local_absmax = GroupReduceMax8(local_absmax);
 
   float y_s = local_absmax / max_8bit;
   y_s = fmaxf(y_s, 1e-10f);
@@ -419,30 +681,8 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
 
   // Quantize and pack into 16 fp8/int8 bytes (= uint4). VEC_SIZE==16 so we
   // fill four 32-bit words, four bytes each.
-  uint32_t packed_lo = 0;
-  uint32_t packed_lo_hi = 0;
-  uint32_t packed_hi_lo = 0;
-  uint32_t packed_hi = 0;
-#pragma unroll
-  for (int i = 0; i < VEC_SIZE; ++i) {
-    float q =
-        fminf(fmaxf(static_cast<float>(regs[i]) * inv_y, min_8bit), max_8bit);
-    DST_DTYPE qb = DST_DTYPE(q);
-    uint8_t byte = *reinterpret_cast<uint8_t*>(&qb);
-    const int shift = (i & 3) * 8;
-    if (i < 4) {
-      packed_lo |= static_cast<uint32_t>(byte) << shift;
-    } else if (i < 8) {
-      packed_lo_hi |= static_cast<uint32_t>(byte) << shift;
-    } else if (i < 12) {
-      packed_hi_lo |= static_cast<uint32_t>(byte) << shift;
-    } else {
-      packed_hi |= static_cast<uint32_t>(byte) << shift;
-    }
-  }
-
-  uint4 packed_out =
-      make_uint4(packed_lo, packed_lo_hi, packed_hi_lo, packed_hi);
+  uint4 packed_out = QuantizeRegisterGroup<T, DST_DTYPE, VEC_SIZE, true>(
+      regs, inv_y, min_8bit, max_8bit);
   DST_DTYPE* group_output =
       static_cast<DST_DTYPE*>(output_q) +
       static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
@@ -473,6 +713,14 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
 
   STD_TORCH_CHECK(input.is_contiguous());
   STD_TORCH_CHECK(output_q.is_contiguous());
+  STD_TORCH_CHECK(
+      Is16ByteAligned(input.data_ptr()),
+      "per_token_group_quant_8bit_packed requires 16-byte aligned input "
+      "data_ptr for uint4 vectorized loads.");
+  STD_TORCH_CHECK(
+      Is16ByteAligned(output_q.data_ptr()),
+      "per_token_group_quant_8bit_packed requires 16-byte aligned output_q "
+      "data_ptr for uint4 vectorized stores.");
 
   const int64_t k = input.size(-1);
   STD_TORCH_CHECK(k % group_size == 0, "input last dim k=", k,

tests/kernels/quantization/test_per_token_group_quant.py

@@ -11,7 +11,17 @@ from vllm.platforms import current_platform
 
 
 @pytest.mark.parametrize(
-    "shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512)]
+    "shape",
+    [
+        (31, 128),
+        (32, 128),
+        (63, 256),
+        (64, 256),
+        (16, 512),
+        (31, 384),
+        (33, 640),
+        (65, 768),
+    ],
 )
 @pytest.mark.parametrize("column_major", [False, True])
 @pytest.mark.parametrize("tma_aligned", [False, True])
@@ -402,7 +412,10 @@ def test_per_token_group_quant_fp8_packed_large_mn():
     assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"
 
 
-@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
+@pytest.mark.parametrize(
+    "shape",
+    [(32, 128), (64, 256), (16, 512), (31, 384), (33, 640), (65, 768)],
+)
 @pytest.mark.parametrize("group_size", [64, 128])
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 def test_per_token_group_quant_int8(shape, group_size: int):