Self-review: clean up stream synchronization

Author: tdene

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -1696,6 +1696,7 @@ async def async_generate_output_tokens_dynamic_batch(
             # Forward pass produces only base logits. When speculative decoding is
             # active, MTP logits are computed serially after verification.
             self._dynamic_step_forward_logits(input_ids, position_ids)
+            torch.cuda.current_stream().wait_event(self._pre_forward_bookkeeping_event)
 
             # Commit Mamba intermediate states before update_requests, which
             # may swap request indices. The Python lists tracking EOS block IDs
@@ -1716,7 +1717,6 @@ async def async_generate_output_tokens_dynamic_batch(
         # NOTE [TDE]: This will be moved once CPU and GPU methods are separated.
         await asyncio.sleep(0)
 
-        torch.cuda.current_stream().wait_event(self._pre_forward_bookkeeping_event)
         with torch.inference_mode():
             return_log_probs, return_top_n_logprobs = self._dynamic_step_log_probs_bookkeeping()
             if self.num_speculative_tokens > 0: