init buffer in kernel

Author: jinzhen-lin

csrc/moe/moe_align_sum_kernels.cu

@@ -222,22 +222,28 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
                           int64_t block_size, torch::Tensor sorted_token_ids,
                           torch::Tensor experts_ids,
                           torch::Tensor num_tokens_post_pad,
-                          torch::Tensor token_cnts_buffer,
-                          torch::Tensor cumsum_buffer) {
+                          bool use_global_memory) {
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   // If we have very large number of experts, we can no longer use shared
   // memory.
   // TODO(simon): the right solution should be calculating the exact right
   // amount of shared memory and use that. The num_experts >= 256 is just a
   // temporary solution to unblock Deepseek V3.
-  if (token_cnts_buffer.numel() > 0 && token_cnts_buffer.numel() > 0) {
+  if (use_global_memory) {
     VLLM_DISPATCH_INTEGRAL_TYPES(
         topk_ids.scalar_type(), "moe_align_block_size_global_mem_kernel", [&] {
           // calc needed amount of shared mem for `tokens_cnts` and `cumsum`
           // tensors
           const int32_t num_thread = max((int32_t)num_experts, WARP_SIZE);
 
+          auto options_int =
+              torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
+          torch::Tensor token_cnts_buffer =
+              torch::empty({(num_experts + 1) * num_experts}, options_int);
+          torch::Tensor cumsum_buffer =
+              torch::empty({num_experts + 1}, options_int);
+
           auto kernel =
               vllm::moe::moe_align_block_size_global_mem_kernel<scalar_t>;
           kernel<<<1, num_thread, 0, stream>>>(

csrc/moe/moe_ops.h

@@ -12,5 +12,4 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
                           int64_t block_size, torch::Tensor sorted_token_ids,
                           torch::Tensor experts_ids,
                           torch::Tensor num_tokens_post_pad,
-                          torch::Tensor token_cnts_buffer,
-                          torch::Tensor cumsum_buffer);
+                          bool use_global_memory);

csrc/moe/torch_bindings.cpp

@@ -20,8 +20,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "                     int block_size, Tensor! sorted_token_ids,"
       "                     Tensor! experts_ids,"
       "                     Tensor! num_tokens_post_pad,"
-      "                     Tensor! token_cnts_buffer,"
-      "                     Tensor! cumsum_buffer) -> ()");
+      "                     bool use_global_memory) -> ()");
   m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
 
 #ifndef USE_ROCM

vllm/_custom_ops.py

@@ -918,17 +918,11 @@ def moe_align_block_size(topk_ids: torch.Tensor, num_experts: int,
                          block_size: int, sorted_token_ids: torch.Tensor,
                          experts_ids: torch.Tensor,
                          num_tokens_post_pad: torch.Tensor,
-                         token_cnts_buffer: Optional[torch.Tensor] = None,
-                         cumsum_buffer: Optional[torch.Tensor] = None,
-                         ) -> None:
-    if token_cnts_buffer is None:
-        token_cnts_buffer = torch.empty((0,), device=topk_ids.device)
-    if cumsum_buffer is None:
-        cumsum_buffer = torch.empty((0,), device=topk_ids.device)
+                         use_global_memory: bool = False) -> None:
     torch.ops._moe_C.moe_align_block_size(topk_ids, num_experts, block_size,
                                           sorted_token_ids, experts_ids,
                                           num_tokens_post_pad,
-                                          token_cnts_buffer, cumsum_buffer)
+                                          use_global_memory)
 
 
 def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,

vllm/config.py

@@ -604,6 +604,7 @@ def _verify_cuda_graph(self) -> None:
                                           self.max_model_len)
 
         if (self.hf_config.model_type == 'deepseek_v3'
+                and self.quantization == "fp8"
                 and not self.enforce_eager):
             logger.warning("CUDA graph is not supported for Deepseek V3 yet, "
                            "fallback to the eager mode.")

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -256,25 +256,10 @@ def moe_align_block_size(
     num_tokens_post_pad = torch.empty((1),
                                       dtype=torch.int32,
                                       device=topk_ids.device)
-    if num_experts >= 256:
-        # For DeepSeek-V3
-        token_cnts_buffer = torch.empty((num_experts + 1) * num_experts,
-                                        dtype=torch.int32,
-                                        device=topk_ids.device)
-        cumsum_buffer = torch.empty(num_experts + 1,
-                                    dtype=torch.int32,
-                                    device=topk_ids.device)
-    else:
-        token_cnts_buffer = torch.empty((0, ),
-                                        dtype=torch.int32,
-                                        device=topk_ids.device)
-        cumsum_buffer = torch.empty((0, ),
-                                    dtype=torch.int32,
-                                    device=topk_ids.device)
-
+    use_global_memory = num_experts >= 256  # for deepseek-v3
     ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
                              expert_ids, num_tokens_post_pad,
-                             token_cnts_buffer, cumsum_buffer)
+                             use_global_memory)
     return sorted_ids, expert_ids, num_tokens_post_pad