fix up

Author: esmeetu

csrc/moe_align_block_size_kernels.cu

@@ -19,7 +19,7 @@ __global__ void moe_align_block_size_kernel(scalar_t *__restrict__ topk_ids,
                                 int32_t num_experts, 
                                 int32_t block_size, 
                                 size_t numel) {
-    const size_t tokens_per_thread = ((numel + blockDim.x - 1) / blockDim.x);
+    const size_t tokens_per_thread = CEILDIV(numel, blockDim.x);
     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];

vllm/model_executor/layers/fused_moe.py

@@ -30,8 +30,6 @@ def fused_moe_kernel(
     stride_bn,
     stride_cm,
     stride_cn,
-    stride_weight,
-    stride_token_id,
     # Meta-parameters
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
@@ -112,7 +110,7 @@ def fused_moe_kernel(
         b_ptrs += BLOCK_SIZE_K * stride_bk
 
     if MUL_ROUTED_WEIGHT:
-        moe_weight = tl.load(topk_weights_ptr + offs_token * stride_weight,
+        moe_weight = tl.load(topk_weights_ptr + offs_token,
                              mask=token_mask,
                              other=0)
         accumulator = accumulator * moe_weight[:, None]
@@ -178,6 +176,10 @@ def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
                             expert_ids: torch.Tensor,
                             num_tokens_post_padded: torch.Tensor,
                             mul_routed_weight: bool, top_k: int, config: dict):
+    
+    assert topk_weights.stride(1) == 1
+    assert sorted_token_ids.stride(0) == 1
+
     grid = lambda META: (triton.cdiv(sorted_token_ids.shape[0], META[
         'BLOCK_SIZE_M']) * triton.cdiv(B.shape[1], META['BLOCK_SIZE_N']), )
 
@@ -200,8 +202,6 @@ def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
         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,

vllm/model_executor/models/deepseek.py

@@ -2,7 +2,7 @@
 # Adapted from
 # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
 # Copyright 2023 The vLLM team.
-# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+# Copyright 2023 DeepSeek-AI and the HuggingFace Inc. team. All rights reserved.
 #
 # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
 # and OPT implementations in this library. It has been modified from its
@@ -22,12 +22,12 @@
 # limitations under the License.
 """Inference-only Deepseek model."""
 from typing import Any, Dict, List, Optional, Tuple
+from transformers import PretrainedConfig
 
 import torch
 import torch.nn.functional as F
 
 from torch import nn
-from vllm.transformers_utils.configs.deepseek import DeepseekConfig
 
 from vllm.model_executor.input_metadata import InputMetadata
 from vllm.model_executor.layers.activation import SiluAndMul
@@ -91,7 +91,7 @@ class DeepseekMoE(nn.Module):
 
     def __init__(
         self,
-        config: DeepseekConfig,
+        config: PretrainedConfig,
         linear_method: Optional[LinearMethodBase] = None,
     ):
         super().__init__()
@@ -264,7 +264,7 @@ class DeepseekDecoderLayer(nn.Module):
 
     def __init__(
         self,
-        config: DeepseekConfig,
+        config: PretrainedConfig,
         layer_idx: int,
         linear_method: Optional[LinearMethodBase] = None,
     ) -> None:
@@ -331,7 +331,7 @@ class DeepseekModel(nn.Module):
 
     def __init__(
         self,
-        config: DeepseekConfig,
+        config: PretrainedConfig,
         linear_method: Optional[LinearMethodBase] = None,
     ) -> None:
         super().__init__()
@@ -372,7 +372,7 @@ class DeepseekForCausalLM(nn.Module):
 
     def __init__(
         self,
-        config: DeepseekConfig,
+        config: PretrainedConfig,
         linear_method: Optional[LinearMethodBase] = None,
     ) -> None:
         super().__init__()

vllm/transformers_utils/configs/__init__.py

@@ -1,7 +1,6 @@
 from vllm.transformers_utils.configs.aquila import AquilaConfig
 from vllm.transformers_utils.configs.baichuan import BaiChuanConfig
 from vllm.transformers_utils.configs.chatglm import ChatGLMConfig
-from vllm.transformers_utils.configs.deepseek import DeepseekConfig
 from vllm.transformers_utils.configs.mpt import MPTConfig
 from vllm.transformers_utils.configs.qwen import QWenConfig
 # RWConfig is for the original tiiuae/falcon-40b(-instruct) and
@@ -14,7 +13,6 @@
     "AquilaConfig",
     "BaiChuanConfig",
     "ChatGLMConfig",
-    "DeepseekConfig",
     "MPTConfig",
     "QWenConfig",
     "RWConfig",