Improve MFU metrics for prefill and verify timing (#29000)

Co-authored-by: Pranjal Shankhdhar <pranjal.ssh@gmail.com>
This commit is contained in:
Lianmin Zheng
2026-06-23 12:26:56 -07:00
committed by GitHub
co-authored by Pranjal Shankhdhar
parent 0c6e8e9477
commit b60185c41c
4 changed files with 46 additions and 10 deletions
@@ -309,6 +309,11 @@ 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)
@@ -337,7 +342,7 @@ class EagerRunner(BaseRunner):
forward_positions = sharded_positions
ctx = (
model_runner.device_timer.wrap(metadata={"category": "extend"})
model_runner.device_timer.wrap(metadata={"category": category})
if model_runner.device_timer
else contextlib.nullcontext()
)