Handle new_block_ids is None. (#533)

Signed-off-by: Qiliang Cui <[email protected]>

Author: QiliangCui

tpu_commons/runner/jax/tpu_jax_runner.py

@@ -1572,11 +1572,13 @@ def _update_states(self, scheduler_output: "VllmSchedulerOutput") -> bool:
             # Update the cached states.
             req_state.num_computed_tokens = num_computed_tokens
             if not resumed_from_preemption:
-                # Append the new blocks to the existing block IDs.
-                for block_ids, new_ids in zip(req_state.block_ids,
-                                              new_block_ids):
-                    block_ids.extend(new_ids)
+                if new_block_ids is not None:
+                    # Append the new blocks to the existing block IDs.
+                    for block_ids, new_ids in zip(req_state.block_ids,
+                                                  new_block_ids):
+                        block_ids.extend(new_ids)
             else:
+                assert new_block_ids is not None
                 # The request is resumed from preemption.
                 # Replace the existing block IDs with the new ones.
                 req_state.block_ids = new_block_ids
@@ -1592,7 +1594,9 @@ def _update_states(self, scheduler_output: "VllmSchedulerOutput") -> bool:
             # Update the persistent batch.
             self.input_batch.num_computed_tokens_cpu[req_index] = (
                 num_computed_tokens)
-            self.input_batch.block_table.append_row(new_block_ids, req_index)
+            if new_block_ids is not None:
+                self.input_batch.block_table.append_row(
+                    new_block_ids, req_index)
 
             # Add spec_token_ids to token_ids_cpu.
             spec_token_ids = (