Add scheduler metrics extension hooks (#29207)

Co-authored-by: Yinghai Lu <yinghai@meta.com>
2026-06-24 15:50:05 -07:00
parent e26bceb81e
commit 7e63feee6f
4 changed files with 52 additions and 10 deletions
@@ -282,8 +282,13 @@ class EagerRunner(BaseRunner):
            kwargs["input_embeds"] = sharded_hidden_states
            forward_positions = sharded_positions

+        category = (
+            "target_verify"
+            if forward_batch.forward_mode.is_target_verify()
+            else "extend"
+        )
        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()
        )