add maybe_run_recv_hook to __exit__

Signed-off-by: Sage Moore <sage@neuralmagic.com>

Author: SageMoore

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -864,13 +864,14 @@ def forward(
                 self.fused_experts.quant_config,
             )
 
-            # assert self.shared_experts is not None
             if self.shared_experts is not None:
                 shared_output = self.shared_experts(a1)
 
+            # If DBO is being used, register the hook with the ubatch context
+            # and call it in dbo_maybe_run_recv_hook instead of passing it to
+            # the receiver.
             dbo_register_recv_hook(hook)
             dbo_yield()
-
             if dbo_enabled():
                 hook = None
 

vllm/v1/worker/ubatching.py

@@ -58,6 +58,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         global _CURRENT_CONTEXTS, _THREAD_ID_TO_CONTEXT
         _CURRENT_CONTEXTS[self.id] = None
         del _THREAD_ID_TO_CONTEXT[threading.get_ident()]
+        self.maybe_run_recv_hook()
         self.cpu_signal_event.set()
         self.cpu_wait_event.clear()
         self.current_stream = self.compute_stream