fix up

Author: zwd003

csrc/moe_alig_block_size_kernels.cu renamed to csrc/moe_align_block_size_kernels.cu

@@ -11,7 +11,7 @@ const static size_t NUM_MAX_EXPERTS = 64;
 
 namespace vllm {
 template <typename scalar_t>
-__global__ void moe_alig_block_size_kernel(scalar_t *__restrict__ topk_ids, 
+__global__ void moe_align_block_size_kernel(scalar_t *__restrict__ topk_ids, 
                                 int32_t *sorted_token_ids, 
                                 int32_t *expert_ids, 
                                 int32_t *total_tokens_post_pad,
@@ -22,7 +22,7 @@ __global__ void moe_alig_block_size_kernel(scalar_t *__restrict__ topk_ids,
     const size_t start_idx = threadIdx.x * tokens_per_thread;
     __shared__ int32_t tokens_cnts[NUM_MAX_EXPERTS + 1][NUM_MAX_EXPERTS];
     __shared__ int32_t cumsum[NUM_MAX_EXPERTS + 1];
-    for(int i = 0;i < num_experts;i++){
+    for(int i = 0; i < num_experts; ++i){
         tokens_cnts[threadIdx.x + 1][i] = 0;
     }
 
@@ -33,23 +33,23 @@ __global__ void moe_alig_block_size_kernel(scalar_t *__restrict__ topk_ids,
     __syncthreads();
 
     tokens_cnts[0][threadIdx.x] = 0;
-    for(int i=1;i<=blockDim.x;++i){
+    for(int i = 1; i <= blockDim.x; ++i){
         tokens_cnts[i][threadIdx.x] += tokens_cnts[i-1][threadIdx.x];
     }
 
     __syncthreads();
 
-    if(threadIdx.x ==0){
+    if(threadIdx.x == 0){
         cumsum[0] = 0;
-        for(int i=1;i<=num_experts;++i){
+        for(int i = 1; i <= num_experts; ++i){
             cumsum[i] = cumsum[i-1] + (tokens_cnts[blockDim.x][i - 1] + block_size - 1) / block_size * block_size;
         }
         *total_tokens_post_pad = cumsum[num_experts];
     }
 
     __syncthreads();
 
-    for(int i= cumsum[threadIdx.x];i<cumsum[threadIdx.x + 1];i += block_size){
+    for(int i = cumsum[threadIdx.x];i < cumsum[threadIdx.x + 1];i += block_size){
         expert_ids[i / block_size] = threadIdx.x;
     }
 
@@ -62,7 +62,7 @@ __global__ void moe_alig_block_size_kernel(scalar_t *__restrict__ topk_ids,
 }
 }
 
-void moe_alig_block_size(
+void moe_align_block_size(
     torch::Tensor topk_ids,
     int num_experts,
     int block_size,
@@ -73,7 +73,7 @@ void moe_alig_block_size(
     assert(num_experts <= NUM_MAX_EXPERTS);
     VLLM_DISPATCH_INTEGRAL_TYPES(
         topk_ids.scalar_type(), "moe_alig_block_size_kernel", [&] {
-        vllm::moe_alig_block_size_kernel<scalar_t><<<1, num_experts, 0, stream>>>(
+        vllm::moe_align_block_size_kernel<scalar_t><<<1, num_experts, 0, stream>>>(
             topk_ids.data_ptr<scalar_t>(), 
             sorted_token_ids.data_ptr<int32_t>(), 
             experts_ids.data_ptr<int32_t>(), 

csrc/ops.h

@@ -122,7 +122,7 @@ void register_graph_buffers(fptr_t _fa, const std::vector<std::string> &handles,
                             const std::vector<std::vector<int64_t>> &offsets);
 #endif
 
-void moe_alig_block_size(
+void moe_align_block_size(
   torch::Tensor topk_ids,
   int num_experts,
   int block_size,

csrc/pybind.cpp

@@ -57,8 +57,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   ops.def("gptq_shuffle", &gptq_shuffle, "Post processing for GPTQ");
   ops.def("squeezellm_gemm", &squeezellm_gemm, "Quantized GEMM for SqueezeLLM");
   ops.def(
-      "moe_alig_block_size",
-      &moe_alig_block_size,
+      "moe_align_block_size",
+      &moe_align_block_size,
       "Aligning the number of tokens to be processed by each expert such that it is divisible by the block size.");
 
   // Cache ops

setup.py

@@ -305,7 +305,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/moe_alig_block_size_kernels.cu",
+    "csrc/moe_align_block_size_kernels.cu",
     "csrc/pybind.cpp",
 ]
 

vllm/model_executor/layers/fused_moe.py

@@ -16,7 +16,6 @@ def fused_moe_kernel(
     expert_ids_ptr,
     num_tokens_post_padded_ptr,
     # Matrix dimensions
-    M,
     N,
     K,
     EM,
@@ -86,10 +85,9 @@ def fused_moe_kernel(
     a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
                       offs_k[None, :] * stride_ak)
 
-    #
-    off_experts = tl.load(expert_ids_ptr + pid_m) * stride_be
-    b_ptrs = b_ptr + off_experts + (offs_k[:, None] * stride_bk +
-                                    offs_bn[None, :] * stride_bn)
+    off_experts = tl.load(expert_ids_ptr + pid_m)
+    b_ptrs = b_ptr + off_experts * stride_be + (offs_k[:, None] * stride_bk +
+                                                offs_bn[None, :] * stride_bn)
 
     # -----------------------------------------------------------
     # Iterate to compute a block of the C matrix.
@@ -129,7 +127,7 @@ def fused_moe_kernel(
     tl.store(c_ptrs, accumulator, mask=c_mask)
 
 
-def alig_block_size(
+def moe_align_block_size(
         topk_ids: torch.Tensor, block_size: int,
         num_experts: int) -> (torch.Tensor, torch.Tensor, torch.Tensor):
     """
@@ -169,11 +167,48 @@ def alig_block_size(
     num_tokens_post_pad = torch.empty((1),
                                       dtype=torch.int32,
                                       device=topk_ids.device)
-    ops.moe_alig_block_size(topk_ids, num_experts, block_size, sorted_ids,
-                            expert_ids, num_tokens_post_pad)
+    ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
+                             expert_ids, num_tokens_post_pad)
     return sorted_ids, expert_ids, num_tokens_post_pad
 
 
+def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
+                            topk_weights: torch.Tensor, topk_ids: torch.Tensor,
+                            sorted_token_ids: torch.Tensor,
+                            expert_ids: torch.Tensor,
+                            num_tokens_post_padded: torch.Tensor,
+                            mul_routed_weight: bool, top_k: int, config: dict):
+    grid = lambda META: (triton.cdiv(sorted_token_ids.shape[0], META[
+        'BLOCK_SIZE_M']) * triton.cdiv(B.shape[1], META['BLOCK_SIZE_N']), )
+
+    fused_moe_kernel[grid](
+        A,
+        B,
+        C,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        B.shape[1],
+        B.shape[2],
+        sorted_token_ids.shape[0],
+        topk_ids.numel(),
+        A.stride(0),
+        A.stride(1),
+        B.stride(0),
+        B.stride(2),
+        B.stride(1),
+        C.stride(1),
+        C.stride(2),
+        topk_weights.stride(1),
+        sorted_token_ids.stride(0),
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        top_k=top_k,
+        compute_type=tl.bfloat16 if A.dtype == torch.bfloat16 else tl.float16,
+        **config,
+    )
+
+
 def fused_moe(hidden_states: torch.Tensor,
               w1: torch.Tensor,
               w2: torch.Tensor,
@@ -196,11 +231,12 @@ def fused_moe(hidden_states: torch.Tensor,
     """
     # Check constraints.
     assert hidden_states.shape[1] == w1.shape[2], "Incompatible dimensions"
-    assert hidden_states.is_contiguous(), "Matrix A must be contiguous"
-    assert w1.is_contiguous(), "Matrix B must be contiguous"
+    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
+    assert w1.is_contiguous(), "Expert weights1 must be contiguous"
+    assert w2.is_contiguous(), "Expert weights2 must be contiguous"
     assert hidden_states.dtype in [torch.float16, torch.bfloat16]
-    M, K = hidden_states.shape
-    E, N, K = w1.shape
+    M, _ = hidden_states.shape
+    E, N, _ = w1.shape
 
     config = {
         'BLOCK_SIZE_M': 64,
@@ -227,73 +263,21 @@ def fused_moe(hidden_states: torch.Tensor,
                                       device=hidden_states.device,
                                       dtype=hidden_states.dtype)
 
-    sorted_token_ids, expert_ids, num_tokens_post_padded = alig_block_size(
+    sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
         topk_ids, config['BLOCK_SIZE_M'], E)
-    # 1D launch kernel where each block gets its own program.
-    grid = lambda META: (triton.cdiv(sorted_token_ids.shape[0], META[
-        'BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
 
-    fused_moe_kernel[grid](
-        hidden_states,
-        w1,
-        intermediate_cache1,
-        topk_weights,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        M,
-        N,
-        K,
-        sorted_token_ids.shape[0],
-        topk_ids.numel(),
-        hidden_states.stride(0),
-        hidden_states.stride(1),
-        w1.stride(0),
-        w1.stride(2),
-        w1.stride(1),
-        intermediate_cache1.stride(1),
-        intermediate_cache1.stride(2),
-        topk_weights.stride(1),
-        sorted_token_ids.stride(0),
-        MUL_ROUTED_WEIGHT=False,
-        top_k=topk_ids.shape[1],
-        compute_type=tl.bfloat16
-        if hidden_states.dtype == torch.bfloat16 else tl.float16,
-        **config,
-    )
+    invoke_fused_moe_kernel(hidden_states, w1, intermediate_cache1,
+                            topk_weights, topk_ids, sorted_token_ids,
+                            expert_ids, num_tokens_post_padded, False,
+                            topk_ids.shape[1], config)
 
     ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
 
-    grid = lambda META: (triton.cdiv(sorted_token_ids.shape[0], META[
-        'BLOCK_SIZE_M']) * triton.cdiv(w2.shape[1], META['BLOCK_SIZE_N']), )
-    fused_moe_kernel[grid](
-        intermediate_cache2,
-        w2,
-        intermediate_cache3,
-        topk_weights,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        M,
-        w2.shape[1],
-        w2.shape[2],
-        sorted_token_ids.shape[0],
-        topk_ids.numel(),
-        intermediate_cache2.stride(0),
-        intermediate_cache2.stride(1),
-        w2.stride(0),
-        w2.stride(2),
-        w2.stride(1),
-        intermediate_cache3.stride(1),
-        intermediate_cache3.stride(2),
-        topk_weights.stride(1),
-        sorted_token_ids.stride(0),
-        MUL_ROUTED_WEIGHT=True,
-        top_k=1,  #
-        compute_type=tl.bfloat16
-        if hidden_states.dtype == torch.bfloat16 else tl.float16,
-        **config,
-    )
+    invoke_fused_moe_kernel(intermediate_cache2, w2, intermediate_cache3,
+                            topk_weights, topk_ids, sorted_token_ids,
+                            expert_ids, num_tokens_post_padded, True, 1,
+                            config)
+
     if inplace:
         return torch.sum(intermediate_cache3.view(*intermediate_cache3.shape),
                          dim=1,

vllm/model_executor/models/deepseek.py

@@ -149,16 +149,6 @@ def pack_params(self):
 
         self.w2 = self.w2.view(len(w2), *w2s[0].shape)
 
-    def fused_moe_infer(self, hidden_states: torch.Tensor,
-                        selected_experts: torch.Tensor,
-                        routing_weights: torch.Tensor) -> torch.Tensor:
-        return fused_moe(hidden_states,
-                         self.w1,
-                         self.w2,
-                         routing_weights,
-                         selected_experts,
-                         inplace=True)
-
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         batch_size, sequence_length, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
@@ -175,9 +165,12 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         if self.config.norm_topk_prob:
             routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
 
-        final_hidden_states = self.fused_moe_infer(hidden_states,
-                                                   selected_experts,
-                                                   routing_weights)
+        final_hidden_states = fused_moe(hidden_states,
+                                        self.w1,
+                                        self.w2,
+                                        routing_weights,
+                                        selected_experts,
+                                        inplace=True)
 
         if self.config.n_shared_experts is not None:
             final_hidden_states = final_hidden_states + shared_output
@@ -290,15 +283,16 @@ def __init__(
             max_position_embeddings=max_position_embeddings,
             linear_method=linear_method,
         )
-        self.mlp = DeepseekMoE(config=config,
-                                           linear_method=linear_method) if (config.n_routed_experts is not None and  \
-                                           layer_idx >= config.first_k_dense_replace and layer_idx % config.moe_layer_freq == 0) \
-                                        else DeepseekMLP(
-                                                hidden_size=config.hidden_size,
-                                                intermediate_size=config.intermediate_size,
-                                                hidden_act=config.hidden_act,
-                                                linear_method=linear_method,
-                                            )
+        if (config.n_routed_experts is not None and  \
+            layer_idx >= config.first_k_dense_replace and layer_idx % config.moe_layer_freq == 0):
+            self.mlp = DeepseekMoE(config=config, linear_method=linear_method)
+        else:
+            self.mlp = DeepseekMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                linear_method=linear_method,
+            )
         self.input_layernorm = RMSNorm(config.hidden_size,
                                        eps=config.rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(config.hidden_size,