make buffer optional

Author: jinzhen-lin

vllm/_custom_ops.py

@@ -918,8 +918,13 @@ 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: torch.Tensor,
-                         cumsum_buffer: torch.Tensor) -> None:
+                         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)
     torch.ops._moe_C.moe_align_block_size(topk_ids, num_experts, block_size,
                                           sorted_token_ids, experts_ids,
                                           num_tokens_post_pad,