Add fused_transpose_split_quant kernel

lshpku · lshpku · commit 452bb11f470b · 2025-05-28T08:48:01.000Z
diff --git a/slm/model_zoo/gpt-3/external_ops/fused_quanted_ops/fused_transpose_split_quant.cu b/slm/model_zoo/gpt-3/external_ops/fused_quanted_ops/fused_transpose_split_quant.cu
@@ -0,0 +1,269 @@
+#include "quant_utils.h"
+
+template <typename T, int VecSize>
+struct __align__(sizeof(T) * VecSize) VecType {
+  T val[VecSize];
+  __host__ __device__ inline T& operator[](size_t i) { return val[i]; }
+  __host__ __device__ inline const T& operator[](size_t i) const {
+    return val[i];
+  }
+};
+
+template <int VecSize>
+__device__ void BlockLoad(const phi::bfloat16* X,
+                          __nv_bfloat16 input[4][4],
+                          size_t K) {
+  for (size_t i = 0; i < 4; i++) {
+    size_t off_m = blockIdx.x * 128 + threadIdx.y + i * 32;
+    size_t off_k = blockIdx.y * 128 + threadIdx.x * VecSize;
+    size_t offset = off_m * K + off_k;
+
+    for (size_t j = 0; j < 4; j += VecSize) {
+      if (off_k + j * 32 < K) {
+        size_t idx = offset + j * 32;
+        using LoadT = VecType<__nv_bfloat16, VecSize>;
+        LoadT data = *reinterpret_cast<const LoadT*>(X + idx);
+        for (int k = 0; k < VecSize; k++) {
+          input[i][j + k] = data[k];
+        }
+      }
+    }
+  }
+}
+
+__device__ void BlockColumnMax(const __nv_bfloat16 input[4][4],
+                               __nv_bfloat16 amax[4],
+                               __nv_bfloat16* shm) {
+  // Reduce [(4), 32, 32, 4] => [32, 32, 4]
+  __nv_bfloat16 warp_max[4];
+  for (int i = 0; i < 4; i++) {
+    for (int j = 0; j < 4; j++) {
+      __nv_bfloat16 t = __habs(input[i][j]);
+      warp_max[j] = i == 0 ? t : __hmax(warp_max[j], t);
+    }
+  }
+
+  // Reduce [(32), 32, 4] => [32, 4]
+  for (int i = 0; i < 4; i++) {
+    shm[threadIdx.y * 128 + i * 32 + threadIdx.x] = warp_max[i];
+  }
+  __syncthreads();
+  for (int offset = 16; offset > 0; offset /= 2) {
+    if (threadIdx.y < offset) {
+      for (int i = 0; i < 4; i++) {
+        shm[threadIdx.y * 128 + i * 32 + threadIdx.x] =
+            __hmax(shm[threadIdx.y * 128 + i * 32 + threadIdx.x],
+                   shm[(threadIdx.y + offset) * 128 + i * 32 + threadIdx.x]);
+      }
+    }
+    __syncthreads();
+  }
+
+  for (int i = 0; i < 4; i++) {
+    amax[i] = shm[i * 32 + threadIdx.x];
+  }
+}
+
+template <typename OutT, int VecSize>
+__device__ void BlockStoreScale(float* scale,
+                                __nv_bfloat16 amax[4],
+                                float scale_inv[4],
+                                size_t K) {
+  float scale_out[4];
+  for (int i = 0; i < 4; i++) {
+    float amax_fp32 = static_cast<float>(amax[i]);
+    scale_inv[i] = ComputeScale<__nv_bfloat16, OutT, true>(amax_fp32, 0.0f);
+    scale_out[i] = __frcp_rn(scale_inv[i]);
+  }
+  if (threadIdx.y == 0) {
+    size_t off_m = blockIdx.x;
+    size_t off_k = blockIdx.y * 128 + threadIdx.x * VecSize;
+    size_t offset = off_m * K + off_k;
+
+    for (size_t j = 0; j < 4; j += VecSize) {
+      if (off_k + j * 32 < K) {
+        size_t idx = offset + j * 32;
+        using StoreT = VecType<float, VecSize>;
+        StoreT data;
+        for (int k = 0; k < VecSize; k++) {
+          data[k] = scale_out[j + k];
+        }
+        *reinterpret_cast<StoreT*>(scale + idx) = data;
+      }
+    }
+  }
+}
+
+template <typename OutT, int VecSize>
+__device__ void BlockStoreOut(OutT* out,
+                              const OutT shm[128][129],
+                              const int64_t* __restrict__ tokens_per_expert,
+                              size_t num_experts,
+                              size_t K) {
+  // Find the current expert_idx
+  size_t idx_m = blockIdx.x * 128 + threadIdx.x * 4;
+  size_t expert_idx = 0;
+  size_t tokens_offset = 0;
+  size_t next_tokens_offset = 0;
+  for (; expert_idx < num_experts; expert_idx++) {
+    next_tokens_offset += tokens_per_expert[expert_idx];
+    if (idx_m >= tokens_offset && idx_m < next_tokens_offset) {
+      break;
+    }
+    tokens_offset = next_tokens_offset;
+  }
+
+  for (size_t i = 0; i < 4; i++) {
+    size_t idx_k = blockIdx.y * 128 + threadIdx.y + i * 32;
+    size_t idx = tokens_offset * K + (idx_m - tokens_offset) +
+                 idx_k * tokens_per_expert[expert_idx];
+
+    if (idx_k < K) {
+      using StoreT = VecType<OutT, VecSize>;
+      StoreT data;
+      for (int j = 0; j < VecSize; j++) {
+        data[j] = shm[i * 32 + threadIdx.y][threadIdx.x * 4 + j];
+      }
+      *reinterpret_cast<StoreT*>(out + idx) = data;
+    }
+  }
+}
+
+template <typename OutT, int VecSize>
+__global__ void __launch_bounds__(1024, 2) FusedTransposeSplitQuantKernel(
+    const phi::bfloat16* __restrict__ X,
+    OutT* __restrict__ out,
+    float* __restrict__ scale,
+    const int64_t* __restrict__ tokens_per_expert,
+    size_t num_experts,
+    size_t K) {
+  __shared__ OutT shm[128][129];
+
+  // Load 128x128 elements from X
+  __nv_bfloat16 input[4][4];
+  BlockLoad<VecSize>(X, input, K);
+
+  // Find the maximum of each 128 elements on the M axis
+  __nv_bfloat16 amax[4];
+  BlockColumnMax(input, amax, reinterpret_cast<__nv_bfloat16*>(shm));
+
+  // Compute scale and scale_inv, then store scale back
+  float scale_inv[4];
+  BlockStoreScale<OutT, VecSize>(scale, amax, scale_inv, K);
+
+  // Scale X and save into shared memory with transposed layout
+  for (int i = 0; i < 4; i++) {
+    for (int j = 0; j < 4; j += VecSize) {
+      for (int k = 0; k < VecSize; k++) {
+        float input_fp32 = static_cast<float>(input[i][j + k]);
+        float output_scaled = input_fp32 * scale_inv[j + k];
+        shm[threadIdx.x * VecSize + j * 32 + k][i * 32 + threadIdx.y] =
+            static_cast<OutT>(output_scaled);
+      }
+    }
+  }
+  __syncthreads();
+
+  // Store 128x128 elements back
+  // Note: out is always 4x vectorizable.
+  BlockStoreOut<OutT, 4>(out, shm, tokens_per_expert, num_experts, K);
+}
+
+/**
+ * Quantize on dim[0] of X, transpose dim[0] and dim[1] of X, then
+ * split the output and scale due to tokens_per_expert.
+ *
+ * Inputs:
+ *   X     : [SUM(M_1...M_N), K], bfloat16
+ *   tokens_per_expert
+ *         : python list of value [M_1, M_2, ..., M_N]
+ *
+ * Outputs:
+ *   out   : [K * M_1 + K * M_2 + ... + K * M_N]
+ *   scale : [SUM(M_1...M_N)/128, K]
+ *
+ * Requirements:
+ *   1) M_i % 128 == 0 for each M_i in tokens_per_expert
+ *   2) K <= 65535 * 128
+ */
+std::vector<paddle::Tensor> fused_transpose_split_quant(
+    const paddle::Tensor& X, const std::vector<int64_t>& tokens_per_expert) {
+  PD_CHECK(X.dtype() == paddle::DataType::BFLOAT16);
+
+  std::vector<int64_t> shape = X.shape();
+  PD_CHECK(shape.size() == 2);
+  const int64_t M = shape[0];
+  const int64_t K = shape[1];
+
+  int64_t sum_tokens = 0;
+  for (size_t i = 0; i < tokens_per_expert.size(); i++) {
+    PADDLE_ENFORCE_EQ(tokens_per_expert[i] % 128,
+                      0,
+                      common::errors::InvalidArgument(
+                          "Each tokens_per_expert must be multiple of 128. "
+                          "However, got tokens_per_expert[%d] = %lld.",
+                          i,
+                          tokens_per_expert[i]));
+    sum_tokens += tokens_per_expert[i];
+  }
+
+  PADDLE_ENFORCE_EQ(
+      sum_tokens,
+      M,
+      common::errors::InvalidArgument(
+          "Sum of tokens_per_expert must be equal to X.shape[0]."));
+  PADDLE_ENFORCE_LE(K,
+                    65535 * 128,
+                    common::errors::InvalidArgument(
+                        "X.shape[1] must be no larger than 65535 * 128."));
+
+  // Allocate for out and scale
+  paddle::Tensor out =
+      paddle::empty({K * M}, paddle::DataType::FLOAT8_E4M3FN, X.place());
+  paddle::Tensor scale =
+      paddle::empty({M / 128, K}, paddle::DataType::FLOAT32, X.place());
+
+  // Skip 0-size
+  if (M == 0 || K == 0) {
+    return {out, scale};
+  }
+
+  // Copy tokens_per_expert to device
+  paddle::Tensor tokens_per_expert_cpu =
+      paddle::empty({static_cast<int64_t>(tokens_per_expert.size())},
+                    paddle::DataType::INT64);
+  std::memcpy(tokens_per_expert_cpu.data(),
+              tokens_per_expert.data(),
+              sizeof(int64_t) * tokens_per_expert.size());
+  paddle::Tensor tokens_per_expert_gpu =
+      tokens_per_expert_cpu.copy_to(X.place(), /* blocking= */ false);
+
+  // Launch kernel
+  dim3 grid(M / 128, (K + 127) / 128);
+  dim3 block(32, 32);
+
+#define LAUNCH_KERNEL(VEC_SIZE) \
+  FusedTransposeSplitQuantKernel<phi::float8_e4m3fn, VEC_SIZE> \
+      <<<grid, block>>>(X.data<phi::bfloat16>(), \
+                        out.data<phi::float8_e4m3fn>(), \
+                        scale.data<float>(), \
+                        tokens_per_expert_gpu.data<int64_t>(), \
+                        tokens_per_expert.size(), \
+                        K);
+  if (K % 4 == 0) {
+    LAUNCH_KERNEL(4);
+  } else if (K % 2 == 0) {
+    LAUNCH_KERNEL(2);
+  } else {
+    LAUNCH_KERNEL(1);
+  }
+#undef LAUNCH_KERNEL
+
+  return {out, scale};
+}
+
+PD_BUILD_OP(fused_transpose_split_quant)
+    .Inputs({"X"})
+    .Outputs({"output", "scale"})
+    .Attrs({"tokens_per_expert: std::vector<int64_t>"})
+    .SetKernelFn(PD_KERNEL(fused_transpose_split_quant));
diff --git a/slm/model_zoo/gpt-3/external_ops/setup_fp8.py b/slm/model_zoo/gpt-3/external_ops/setup_fp8.py
@@ -42,6 +42,7 @@ def setup_fused_quant_ops():
                 "fused_quanted_ops/fused_act_dequant_transpose_act_quant.cu",
                 "fused_quanted_ops/fused_spaq.cu",
                 "fused_quanted_ops/fused_stack_transpose_quant.cu",
+                "fused_quanted_ops/fused_transpose_split_quant.cu",
             ],
             extra_compile_args={
                 "cxx": ["-O3", "-w", "-Wno-abi", "-fPIC", "-std=c++17"],
diff --git a/tests/ops/test_fused_transpose_split_quant.py b/tests/ops/test_fused_transpose_split_quant.py
@@ -0,0 +1,49 @@
+import FusedQuantOps as FQO
+import numpy as np
+
+import paddle
+
+
+def restore_transpose_split_quant(out, scale):
+    out = [t.astype('float32') for t in out]
+    out = paddle.concat(out, axis=1).transpose([1, 0])
+    scale = paddle.concat(scale, axis=0)
+    scale = paddle.repeat_interleave(scale, repeats=128, axis=0)
+    return out * scale
+
+
+def run():
+    tokens_per_expert = [24*128, 50*128, 1*128, 128*128, 13*128]
+
+    for seq_len in [1, 127, 2562, 4001, 7168]:
+        print(tokens_per_expert, seq_len)
+
+        x = paddle.randn([sum(tokens_per_expert), seq_len], dtype='bfloat16')
+        x = paddle.clip(x, min=-50, max=50)
+
+        out_raw, scale_raw = FQO.fused_transpose_split_quant(
+            x, tokens_per_expert
+        )
+
+        out, scale = [], []
+        token_offset = 0
+        for tokens in tokens_per_expert:
+            out_offset = seq_len * token_offset
+            out_size = seq_len * tokens
+            out.append(
+                out_raw[out_offset : out_offset + out_size]
+                    .reshape([seq_len, tokens])
+            )
+            scale.append(
+                scale_raw[token_offset // 128 : (token_offset + tokens) // 128]
+            )
+            token_offset += tokens
+
+        x_restore = restore_transpose_split_quant(out, scale)
+        x_cast = x.astype('float32')
+
+        np.testing.assert_allclose(x_cast, x_restore, rtol=0.01, atol=0.3)
+
+
+if __name__ == '__main__':
+    run()
