[Feature] console print statistical metrics (#6339)

* [Feature] console print statistical data

* [Feature] console print statistical data v2 dp_rank

* [Feature] console print statistical data v2 unittest

* [Feature] console print statistical data v3 unittest
This commit is contained in:
CSWYF3634076
2026-02-05 19:20:36 +08:00
committed by GitHub
parent de02a909c8
commit 1c0a2b055f
7 changed files with 266 additions and 0 deletions
@@ -146,6 +146,7 @@ class TestTokenProcessorProcessBatchOutput(unittest.TestCase):
processor.engine_worker_queue = Mock()
processor.split_connector = Mock()
processor.resource_manager = MockResourceManager()
processor.scheduler_metrics_logger = None
task1 = MockTask()
task2 = MockTask()
processor.resource_manager.tasks_list = [task1, task2]
+15
View File
@@ -199,6 +199,21 @@ def test_init_allocates_expected_buffers():
processor_logprob, _, _, _ = _make_processor(enable_logprob=True)
assert list(processor_logprob.output_scores.shape) == [MAX_BSZ * (K + 1), 1]
def test_is_decode_stage():
processor, _, _, _ = _make_processor()
assert processor._is_decode_stage(None) is False
task = types.SimpleNamespace(need_prefill_tokens=None, num_computed_tokens=0)
assert processor._is_decode_stage(task) is False
task = types.SimpleNamespace(need_prefill_tokens=4, num_computed_tokens=3)
assert processor._is_decode_stage(task) is False
task = types.SimpleNamespace(need_prefill_tokens=4, num_computed_tokens=4)
assert processor._is_decode_stage(task) is True
processor_spec, _, _, _ = _make_processor(speculative_method="mtp", enable_logprob=False)
assert processor_spec.output_tokens.shape[0] == SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2