tensor parallel MOE implementation

benchmarks/benchmark_latency.py

@@ -24,6 +24,7 @@ def main(args: argparse.Namespace):
  trust_remote_code=args.trust_remote_code,
  dtype=args.dtype,
  enforce_eager=args.enforce_eager,
+ load_format="dummy",
  )
 
  sampling_params = SamplingParams(

csrc/bincount.cu

@@ -0,0 +1,32 @@
+#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);
+ }
+}
+}
+
+void vllm_bincount(torch::Tensor src, torch::Tensor out) {
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ AT_DISPATCH_ALL_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,
+ "Gather key and value from the cache into contiguous QKV tensors");
 
  // Cache ops
  pybind11::module cache_ops = m.def_submodule("cache_ops", "vLLM cache ops");

setup.py

@@ -213,6 +213,7 @@ def get_torch_arch_list() -> Set[str]:
 ext_modules = []
 
 vllm_extension_sources = [
+ "csrc/bincount.cu",
  "csrc/cache_kernels.cu",
  "csrc/attention/attention_kernels.cu",
  "csrc/pos_encoding_kernels.cu",

tests/kernels/test_moe_grouped_matmul.py

@@ -0,0 +1,62 @@
+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.float32, 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)
+ diff = torch.abs(ref_output - output)
+ mean = torch.mean(diff)
+ max = torch.max(diff)
+ print(f"{mean=}, {max=}")
+
+ torch.allclose(output, ref_output, atol=1e-2, rtol=0)

vllm/model_executor/layers/moe.py

@@ -0,0 +1,256 @@
+from typing import Tuple
+
+import torch
+
+from torch import nn
+import torch.nn.functional as F
+
+import triton
+import triton.language as tl
+
+from vllm._C import ops
+from vllm.model_executor.layers.linear import ReplicatedLinear
+
+from vllm.model_executor.parallel_utils.communication_op import (
+ tensor_model_parallel_all_reduce)
+
+
+class MoE(nn.Module):
+
+ def __init__(
+ self,
+ num_experts: int,
+ top_k: int,
+ hidden_size: int,
+ intermediate_size: int,
+ tp_size: int,
+ ):
+ super().__init__()
+ self.num_total_experts = num_experts
+ self.top_k = top_k
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size // tp_size
+
+ self.gate = ReplicatedLinear(self.hidden_size,
+ self.num_total_experts,
+ bias=False,
+ linear_method=None)
+
+ self.w1s = nn.Parameter(
+ torch.empty(self.num_total_experts, self.hidden_size,
+ self.intermediate_size))
+ self.w2s = nn.Parameter(
+ torch.empty(self.num_total_experts, self.intermediate_size,
+ self.hidden_size))
+ self.w3s = nn.Parameter(
+ torch.empty(self.num_total_experts, self.hidden_size,
+ self.intermediate_size))
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ batch_size, sequence_length, hidden_size = hidden_states.shape
+ hidden_states = hidden_states.view(-1, self.hidden_size)
+ # router_logits: (batch * sequence_length, n_experts)
+ router_logits, _ = self.gate(hidden_states)
+
+ routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+ routing_weights, selected_experts = torch.topk(routing_weights,
+ self.top_k,
+ dim=-1)
+ routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+
+ expanded_hidden_states, experts_range, expanded_weights, experts_indices = \
+ self.expand_and_permutate_hidden_states(
+ hidden_states, selected_experts, routing_weights)
+
+ expanded_hidden_states = self.grouped_mlp(expanded_hidden_states,
+ experts_range, self.w1s.data,
+ self.w2s.data, self.w3s.data)
+
+ expanded_hidden_states.mul_(expanded_weights.unsqueeze(-1))
+
+ tensor_model_parallel_all_reduce(expanded_hidden_states)
+
+ return self.merge_expert_outputs(expanded_hidden_states,
+ experts_indices).view(
+ batch_size, sequence_length,
+ hidden_size)
+
+ def expand_and_permutate_hidden_states(
+ self,
+ hidden_states: torch.Tensor, # [batch_size, hidden_size]
+ selected_experts: torch.Tensor, # [batch_size, top_k_experts]
+ routing_weights: torch.Tensor, # [batch_size, top_k_experts]
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ _, experts_indices = torch.sort(selected_experts.view(-1), dim=-1)
+ cum_experts_range = torch.zeros(self.num_total_experts + 1,
+ dtype=torch.int32,
+ device=hidden_states.device)
+ num_rows_per_expert = torch.zeros(self.num_total_experts,
+ dtype=torch.int32,
+ device=hidden_states.device)
+ ops.bincount(selected_experts.view(-1), num_rows_per_expert)
+ torch.cumsum(num_rows_per_expert, dim=0, out=cum_experts_range[1:])
+ expanded_weights = routing_weights.view(-1)[experts_indices]
+ return hidden_states[experts_indices.div_(
+ self.top_k, rounding_mode="floor"
+ )], cum_experts_range, expanded_weights, experts_indices
+
+ def grouped_mlp(
+ self,
+ expanded_hidden_states: torch.
+ Tensor, # [batch_size * top_k_experts, hidden_size]
+ cum_experts_range: torch.Tensor, # [num_experts + 1]
+ w1s: torch.Tensor, # [num_experts, hidden_size, ffn_dim]
+ w2s: torch.Tensor, # [num_experts, ffn_dim, hidden_size]
+ w3s: torch.Tensor, # [num_experts, hidden_size, ffn_dim]
+ ) -> torch.Tensor: # [batch_size * top_k_experts, hidden_size]
+ grouped_w1_out = grouped_matmul(expanded_hidden_states,
+ cum_experts_range, w1s, "silu")
+ grouped_w3_out = grouped_matmul(expanded_hidden_states,
+ cum_experts_range, w3s)
+ grouped_w1_out.mul_(grouped_w3_out)
+ return grouped_matmul(grouped_w1_out, cum_experts_range, w2s)
+
+ def merge_expert_outputs(
+ self,
+ expanded_hidden_states: torch.
+ Tensor, # [batch_size * top_k_experts, hidden_size]
+ expert_indicies, # [batch_size * top_k_experts]
+ ) -> torch.Tensor:
+ out = torch.zeros(expanded_hidden_states.shape[0] // self.top_k,
+ self.hidden_size,
+ device=expanded_hidden_states.device,
+ dtype=expanded_hidden_states.dtype)
+ out.index_add_(0, expert_indicies, expanded_hidden_states)
+ return out
+
+
+@triton.jit
+def grouped_matmul_kernel(
+ # device tensor of matrices pointers
+ fused_input_ptr,
+ cum_input_group_range,
+ fused_b_ptr,
+ fused_output_ptr,
+ group_size,
+ n,
+ k,
+ lda,
+ ldb,
+ ldc,
+ # number of virtual SM
+ NUM_SM: tl.constexpr,
+ # tile sizes
+ BLOCK_SIZE_M: tl.constexpr,
+ BLOCK_SIZE_N: tl.constexpr,
+ BLOCK_SIZE_K: tl.constexpr,
+ ACTIVATION: tl.constexpr,
+):
+ tile_idx = tl.program_id(0)
+ last_problem_end = 0
+ for g in range(group_size):
+ # get the gemm size of the current problem
+ a_offset = tl.load(cum_input_group_range + g)
+ gm = tl.load(cum_input_group_range + g + 1) - a_offset
+ gn = n
+ gk = k
+ num_m_tiles = tl.cdiv(gm, BLOCK_SIZE_M)
+ num_n_tiles = tl.cdiv(gn, BLOCK_SIZE_N)
+ num_tiles = num_m_tiles * num_n_tiles
+ # iterate through the tiles in the current gemm problem
+ while (tile_idx >= last_problem_end
+ and tile_idx < last_problem_end + num_tiles):
+
+ # pick up a tile from the current gemm problem
+ k = gk
+ a_ptr = fused_input_ptr + a_offset * lda
+ b_ptr = fused_b_ptr + g * k * n
+ c_ptr = fused_output_ptr + a_offset * ldc
+ # figure out tile coordinates
+ tile_idx_in_gemm = tile_idx - last_problem_end
+ tile_m_idx = tile_idx_in_gemm // num_n_tiles
+ tile_n_idx = tile_idx_in_gemm % num_n_tiles
+
+ # do regular gemm here
+ offs_am = tile_m_idx * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+ offs_bn = tile_n_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+ offs_k = tl.arange(0, BLOCK_SIZE_K)
+ a_ptrs = a_ptr + offs_am[:, None] * lda + offs_k[None, :]
+ b_ptrs = b_ptr + offs_k[:, None] * ldb + offs_bn[None, :]
+ accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N),
+ dtype=tl.float32)
+ for kk in range(0, tl.cdiv(k, BLOCK_SIZE_K)):
+ # hint to Triton compiler to do proper loop pipelining
+ tl.multiple_of(a_ptrs, [16, 16])
+ tl.multiple_of(b_ptrs, [16, 16])
+
+ a = tl.load(a_ptrs,
+ mask=offs_k[None, :] < k - kk * BLOCK_SIZE_K,
+ other=0.0)
+ b = tl.load(b_ptrs,
+ mask=offs_k[:, None] < k - kk * BLOCK_SIZE_K,
+ other=0.0)
+ accumulator += tl.dot(a, b)
+ a_ptrs += BLOCK_SIZE_K
+ b_ptrs += BLOCK_SIZE_K * ldb
+
+ if ACTIVATION == "silu":
+ accumulator = silu(accumulator)
+ c = accumulator.to(tl.float16)
+
+ offs_cm = tile_m_idx * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+ offs_cn = tile_n_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+ c_ptrs = c_ptr + ldc * offs_cm[:, None] + offs_cn[None, :]
+ c_mask = (offs_cm[:, None] < gm) & (offs_cn[None, :] < gn)
+
+ tl.store(c_ptrs, c, mask=c_mask)
+
+ # go to the next tile by advancing NUM_SM
+ tile_idx += NUM_SM
+
+ # get ready to go to the next gemm problem
+ last_problem_end = last_problem_end + num_tiles
+
+
+@triton.jit
+def silu(x):
+ return x * tl.sigmoid(x)
+
+
+def grouped_matmul(fused_input: torch.Tensor,
+ cum_group_range: torch.Tensor,
+ fused_group_b: torch.Tensor,
+ activation: str = ""):
+ device = torch.device('cuda')
+ assert cum_group_range.shape[0] == fused_group_b.shape[0] + 1
+ group_size = cum_group_range.shape[0] - 1
+ output = torch.zeros(fused_input.shape[0],
+ fused_group_b.shape[2],
+ device=device,
+ dtype=fused_input.dtype)
+ BLOCK_SIZE_N = 64
+ num_warps = 2
+ if fused_input.shape[0] >= 8:
+ num_warps = 4
+ BLOCK_SIZE_N = 128
+ # we use a fixed number of CTA, and it's auto-tunable
+ grid = lambda META: (META['NUM_SM'], )
+ grouped_matmul_kernel[grid](fused_input,
+ cum_group_range,
+ fused_group_b,
+ output,
+ group_size,
+ n=fused_group_b.shape[2],
+ k=fused_group_b.shape[1],
+ lda=fused_input.stride(0),
+ ldb=fused_group_b.stride(1),
+ ldc=output.stride(0),
+ ACTIVATION=activation,
+ BLOCK_SIZE_M=16,
+ BLOCK_SIZE_N=BLOCK_SIZE_N,
+ BLOCK_SIZE_K=32,
+ NUM_SM=128,
+ num_warps=num_warps,
+ num_stages=5),
+
+ return output