[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
+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