Revert "Improve MFU metrics for prefill and verify timing" (#29079)

This commit is contained in:
Lianmin Zheng
2026-06-23 15:46:22 -07:00
committed by GitHub
parent 11e7c9e0e6
commit ecab3f322e
4 changed files with 10 additions and 46 deletions
@@ -3051,16 +3051,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
):
# Prefill cuda graph (piecewise).
kwargs = self._extend_forward_kwargs(forward_batch, pp_proxy_tensors)
category = (
"target_verify"
if forward_batch.forward_mode.is_target_verify()
else "extend"
)
# TODO: device_timer.wrap is too broad here — it also includes
# load_batch time. Move timing into the prefill cuda graph runner
# to capture only the model.forward part.
ctx = (
self.device_timer.wrap(metadata={"category": category})
self.device_timer.wrap(metadata={"category": "extend"})
if self.device_timer
else contextlib.nullcontext()
)
@@ -250,11 +250,6 @@ class EagerRunner(BaseRunner):
) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
model_runner = self.model_runner
kwargs = model_runner._extend_forward_kwargs(forward_batch, pp_proxy_tensors)
category = (
"target_verify"
if forward_batch.forward_mode.is_target_verify()
else "extend"
)
if not model_runner.server_args.enable_pdmux:
forward_batch = self.load_batch(forward_batch, pp_proxy_tensors)
@@ -283,7 +278,7 @@ class EagerRunner(BaseRunner):
forward_positions = sharded_positions
ctx = (
model_runner.device_timer.wrap(metadata={"category": category})
model_runner.device_timer.wrap(metadata={"category": "extend"})
if model_runner.device_timer
else contextlib.nullcontext()
)