[Bugfix] Fix TP > 1 for new granite (vllm-project#8544)

Signed-off-by: Joe Runde <[email protected]>
Signed-off-by: LeiWang1999 <[email protected]>

Author: joerunde

vllm/model_executor/models/granite.py

@@ -428,7 +428,8 @@ def compute_logits(
             sampling_metadata: SamplingMetadata) -> Optional[torch.Tensor]:
         logits = self.logits_processor(self.lm_head, hidden_states,
                                        sampling_metadata)
-        logits /= self.config.logits_scaling
+        if logits is not None:
+            logits /= self.config.logits_scaling
         return logits
 
     def sample(