[1.21 cherry-pick] Fix async callback ordering (#1023) (#1028)

Cherry-pick of #1023

Co-authored-by: Michał Kuligowski <mkuligowski@habana.ai>

Author: madamczyk-intel

vllm/worker/hpu_model_runner.py

@@ -2725,6 +2725,8 @@ def try_revert_dummy_output_tokens():
 
                 if use_delayed_sampling:
                     fake_output = self._delayed_sampler_outputs(model_input)
+                elif model_input.async_callback is not None:
+                    model_input.async_callback()
 
                 with self.profiler.record_event(
                         'internal', ('sample_'
@@ -2746,7 +2748,8 @@ def try_revert_dummy_output_tokens():
                         self.cached_step_outputs.append(output)
                         self.cached_step_inputs.append(model_input)
                 htorch.core.mark_step()
-                if model_input.async_callback is not None:
+                if use_delayed_sampling \
+                   and model_input.async_callback is not None:
                     model_input.async_callback()
                 if i < num_steps - 1:
                     if i == 0: