profile decode path and optimize tiny batch moe

vllm/worker/model_runner.py

@@ -2,9 +2,11 @@ import dataclasses
 import gc
 import inspect
 import itertools
+import os
 import time
 import warnings
 import weakref
+from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set,
                     Tuple, Type, TypeVar, Union)
@@ -62,6 +64,56 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+_ENGINEX_PROFILE_ENABLED = os.getenv("ENGINEX_PROFILE_DECODE", "0") == "1"
+_ENGINEX_PROFILE_EVERY = int(os.getenv("ENGINEX_PROFILE_EVERY", "32"))
+_ENGINEX_PROFILE_SYNC = os.getenv("ENGINEX_PROFILE_SYNC", "1") != "0"
+_enginex_profile_stats: Dict[str, List[float]] = {}
+_enginex_profile_steps = 0
+
+
+def _enginex_profile_active() -> bool:
+    return _ENGINEX_PROFILE_ENABLED and torch.cuda.is_available()
+
+
+def _enginex_profile_sync() -> None:
+    if _ENGINEX_PROFILE_SYNC:
+        torch.cuda.synchronize()
+
+
+@contextmanager
+def _enginex_profile(label: str):
+    if not _enginex_profile_active():
+        yield
+        return
+    _enginex_profile_sync()
+    start = time.perf_counter()
+    try:
+        yield
+    finally:
+        _enginex_profile_sync()
+        elapsed_ms = (time.perf_counter() - start) * 1000.0
+        stat = _enginex_profile_stats.setdefault(label, [0.0, 0.0])
+        stat[0] += elapsed_ms
+        stat[1] += 1.0
+
+
+def _enginex_profile_log(mode: Optional[str]) -> None:
+    global _enginex_profile_steps
+    if not _enginex_profile_active():
+        return
+    _enginex_profile_steps += 1
+    if _enginex_profile_steps % max(_ENGINEX_PROFILE_EVERY, 1) != 0:
+        return
+    parts = []
+    for label, (total_ms, count) in sorted(
+            _enginex_profile_stats.items(),
+            key=lambda item: item[1][0],
+            reverse=True):
+        avg_ms = total_ms / max(count, 1.0)
+        parts.append(f"{label}: total={total_ms:.2f}ms avg={avg_ms:.3f}ms n={int(count)}")
+    logger.info("[ENGINEX_PROFILE_MODEL_RUNNER] steps=%d mode=%s %s",
+                _enginex_profile_steps, mode, " | ".join(parts))
+
 LORA_WARMUP_RANK = 8
 _BATCH_SIZE_ALIGNMENT = 8
 # all the token sizes that **can** be captured by cudagraph.
@@ -1633,7 +1685,9 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
                 model_input.prompt_adapter_requests,
                 model_input.prompt_adapter_mapping)
 
-        self.attn_state.begin_forward(model_input)
+        profile_mode = "prompt" if model_input.is_prompt else "decode"
+        with _enginex_profile(f"{profile_mode}.attn_begin_forward"):
+            self.attn_state.begin_forward(model_input)
 
         # Currently cuda graph is only supported by the decode phase.
         assert model_input.attn_metadata is not None
@@ -1661,16 +1715,17 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
             model_forward_end = torch.cuda.Event(enable_timing=True)
             model_forward_start.record()
 
-        with set_forward_context(model_input.attn_metadata):
-            hidden_or_intermediate_states = model_executable(
-                input_ids=model_input.input_tokens,
-                positions=model_input.input_positions,
-                kv_caches=kv_caches,
-                attn_metadata=model_input.attn_metadata,
-                intermediate_tensors=intermediate_tensors,
-                **MultiModalInputs.as_kwargs(multi_modal_kwargs,
-                                             device=self.device),
-                **seqlen_agnostic_kwargs)
+        with _enginex_profile(f"{profile_mode}.model_forward"):
+            with set_forward_context(model_input.attn_metadata):
+                hidden_or_intermediate_states = model_executable(
+                    input_ids=model_input.input_tokens,
+                    positions=model_input.input_positions,
+                    kv_caches=kv_caches,
+                    attn_metadata=model_input.attn_metadata,
+                    intermediate_tensors=intermediate_tensors,
+                    **MultiModalInputs.as_kwargs(multi_modal_kwargs,
+                                                 device=self.device),
+                    **seqlen_agnostic_kwargs)
 
         if (self.observability_config is not None
                 and self.observability_config.collect_model_forward_time):
@@ -1695,20 +1750,23 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
                     torch.tensor(model_forward_time + orig_model_forward_time))
             return hidden_or_intermediate_states
 
-        logits = self.model.compute_logits(hidden_or_intermediate_states,
-                                           model_input.sampling_metadata)
+        with _enginex_profile(f"{profile_mode}.compute_logits"):
+            logits = self.model.compute_logits(hidden_or_intermediate_states,
+                                               model_input.sampling_metadata)
 
         if not self.is_driver_worker:
             return []
 
         if model_input.async_callback is not None:
-            model_input.async_callback()
+            with _enginex_profile(f"{profile_mode}.async_callback"):
+                model_input.async_callback()
 
         # Sample the next token.
-        output: SamplerOutput = self.model.sample(
-            logits=logits,
-            sampling_metadata=model_input.sampling_metadata,
-        )
+        with _enginex_profile(f"{profile_mode}.sample"):
+            output: SamplerOutput = self.model.sample(
+                logits=logits,
+                sampling_metadata=model_input.sampling_metadata,
+            )
         if (self.observability_config is not None
                 and self.observability_config.collect_model_forward_time
                 and output is not None):
@@ -1741,6 +1799,7 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
 
             output.hidden_states = hidden_states
 
+        _enginex_profile_log(profile_mode)
         return [output]