lint

Author: khatwanimohit

src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py

@@ -82,7 +82,10 @@ def generate_maxtext_config(vllm_config: VllmConfig, mesh: Mesh) -> pyconfig.Hyp
   else:
     max_tp_size = mesh.shape[ShardingAxisName.ATTN_HEAD]
 
-  if vllm_config.model_config.get_total_num_kv_heads() < max_tp_size:
+  if (
+      max_tp_size % vllm_config.model_config.get_total_num_kv_heads() == 0
+      and vllm_config.model_config.get_total_num_kv_heads() < max_tp_size
+  ):
     max_logging.log(
         f"Padding num_kv_heads from {vllm_config.model_config.get_total_num_kv_heads()} to {max_tp_size} to match tp_size."
     )