tensor parallel MOE implementation

csrc/dispatch_utils.h

@@ -14,3 +14,14 @@
 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
  AT_DISPATCH_SWITCH( \
  TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
+
+#define VLLM_DISPATCH_CASE_INTEGRAL_TYPES(...) \
+ AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \
+ AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \
+ AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \
+ AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \
+ AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__)
+
+#define VLLM_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \
+ AT_DISPATCH_SWITCH( \
+ TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__))

csrc/misc_kernels.cu

@@ -0,0 +1,35 @@
+// code adapted from https://github.com/rusty1s/pytorch_bincount
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <ATen/ATen.h>
+#include <THC/THCAtomics.cuh>
+
+#include "cuda_compat.h"
+#include "dispatch_utils.h"
+
+#define THREADS 1024
+#define BLOCKS(N) (N + THREADS - 1) / THREADS
+
+namespace vllm {
+template <typename scalar_t>
+__global__ void bincount_kernel(scalar_t *__restrict__ src, int32_t *out,
+ size_t numel) {
+ const size_t index = blockIdx.x * blockDim.x + threadIdx.x;
+ const size_t stride = blockDim.x * gridDim.x;
+ for (ptrdiff_t i = index; i < numel; i += stride) {
+ atomicAdd(out + (ptrdiff_t)src[i], 1);
+ }
+}
+}
+
+// create a custom bincount since pytorch's bincount is
+// not cudagraph capturable.
+void vllm_bincount(torch::Tensor src, torch::Tensor out) {
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ VLLM_DISPATCH_INTEGRAL_TYPES(
+ src.scalar_type(), "bincount_kernel", [&] {
+ vllm::bincount_kernel<scalar_t><<<BLOCKS(src.numel()), THREADS, 0, stream>>>(
+ src.data<scalar_t>(), out.data<int32_t>(), src.numel());
+ });
+}

csrc/ops.h

@@ -89,3 +89,7 @@ torch::Tensor gptq_gemm(
 void gptq_shuffle(
  torch::Tensor q_weight,
  torch::Tensor q_perm);
+
+void vllm_bincount(
+ torch::Tensor src,
+ torch::Tensor out);

csrc/pybind.cpp

@@ -55,6 +55,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  ops.def("gptq_gemm", &gptq_gemm, "Quantized GEMM for GPTQ");
  ops.def("gptq_shuffle", &gptq_shuffle, "Post processing for GPTQ");
  ops.def("squeezellm_gemm", &squeezellm_gemm, "Quantized GEMM for SqueezeLLM");
+ ops.def(
+ "bincount",
+ &vllm_bincount,
+ "cuda-graph compatible bincount implementation");
 
  // Cache ops
  pybind11::module cache_ops = m.def_submodule("cache_ops", "vLLM cache ops");

setup.py

@@ -221,6 +221,7 @@ def get_torch_arch_list() -> Set[str]:
  "csrc/quantization/squeezellm/quant_cuda_kernel.cu",
  "csrc/quantization/gptq/q_gemm.cu",
  "csrc/cuda_utils_kernels.cu",
+ "csrc/misc_kernels.cu",
  "csrc/pybind.cpp",
 ]
 

tests/kernels/test_moe_grouped_matmul.py

@@ -0,0 +1,57 @@
+import itertools
+import random
+import pytest
+import torch
+
+from vllm.model_executor.layers.moe import grouped_matmul
+
+
+def ref_grouped_matmul(
+ fused_input: torch.Tensor,
+ cum_group_range: torch.Tensor,
+ fused_group_b: torch.Tensor,
+ activation: str = "",
+) -> torch.Tensor:
+ num_groups = cum_group_range.shape[0] - 1
+ output = torch.zeros(fused_input.shape[0],
+ fused_group_b.shape[2],
+ device=fused_input.device,
+ dtype=fused_input.dtype)
+ for i in range(num_groups):
+ group_i = fused_input[cum_group_range[i]:cum_group_range[i + 1]]
+ group_i_b = fused_group_b[i]
+ output[cum_group_range[i]:cum_group_range[i + 1]] = group_i @ group_i_b
+
+ if activation == "silu":
+ output = torch.nn.functional.silu(output)
+ return output
+
+
+@pytest.mark.parametrize("group_size", [1, 2, 4, 8, 16])
+@pytest.mark.parametrize("m", [1, 5, 33, 81])
+@pytest.mark.parametrize("n", [128, 1024, 2000])
+@pytest.mark.parametrize("k", [128, 1024, 2000])
+@pytest.mark.parametrize("activation", ["", "silu"])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_moe_grouped_matmul(
+ group_size: int,
+ m: int,
+ n: int,
+ k: int,
+ activation: str,
+ dtype: torch.dtype,
+):
+ groups = [random.randint(1, m) for _ in range(group_size)]
+ batch_size = sum(groups)
+ fused_input = torch.randn(batch_size, k, dtype=dtype, device="cuda")
+ cum_group_range = torch.tensor([0] + list(itertools.accumulate(groups)),
+ dtype=torch.int32,
+ device="cuda")
+ fused_group_b = torch.randn(group_size, k, n, dtype=dtype, device="cuda")
+
+ ref_output = ref_grouped_matmul(fused_input, cum_group_range,
+ fused_group_b, activation)
+
+ output = grouped_matmul(fused_input, cum_group_range, fused_group_b,
+ activation)
+ assert torch.allclose(output, ref_output, rtol=0.01, atol=1e-3)