make sure num tokens divisible by tp_size

Signed-off-by: Wuxun Zhang <[email protected]>

Author: wuxun-zhang

vllm_gaudi/v1/worker/hpu_dp_utils.py

@@ -22,6 +22,11 @@ def make(
         dp_size = vllm_config.parallel_config.data_parallel_size
         tp_size = vllm_config.parallel_config.tensor_parallel_size
 
+        if num_tokens % tp_size != 0:
+            # make sure num_tokens is enough to be divided by tp_size for
+            # sequence parallel MOE
+            num_tokens = (num_tokens // tp_size + 1) * tp_size
+
         num_tokens_across_dp = num_tokens * dp_size
 
         dtype = vllm_config.model_config.dtype