Review compliance 2

Author: remi-or

src/transformers/generation/continuous_batching/cache.py

@@ -166,7 +166,8 @@ def __init__(
 
         # Check if the KV heads are part of the TP plan. If they are not, the cache does not need plan for TP.
         # TODO: this is fragile. If your model fails to TP properly because of this, please open an issue.
-        kv_is_tp = "layers.*.self_attn.k_proj" in config.tp_plan and "layers.*.self_attn.v_proj" in config.tp_plan
+        tp_plan = getattr(config, "tp_plan", {})
+        kv_is_tp = "layers.*.self_attn.k_proj" in tp_plan and "layers.*.self_attn.v_proj" in tp_plan
 
         # If the KV heads are TP'ed, each KV head is dispatched to a different GPU, so the effective number of KV heads
         # per GPU is simply divided by the TP size

src/transformers/generation/continuous_batching/distributed.py

@@ -60,7 +60,6 @@ def __init__(self, device_mesh: DeviceMesh | None) -> None:
         self.dp_size = self.world_size // self.tp_size
 
     def infer_if_tp_driver(self) -> bool:
-        """Infer if the current process is the TP driver."""
         return self.tp_local_rank == 0
 
     def destroy_cpu_comm_group(self) -> None: