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)