11// Copyright (c) MLLM Team.
22// Licensed under the MIT License.
33
4+ #include < string>
5+ #include < fstream>
46#include < vector>
57#include < sstream>
68#include < thread>
79#include < chrono>
10+ #include < algorithm> // For std::transform
811
912#include < mllm/mllm.hpp>
1013#include < mllm/utils/Argparse.hpp>
1619
1720#include " models/All.hpp"
1821
22+ #ifndef MLLM_GIT_COMMIT_HASH
23+ #define MLLM_GIT_COMMIT_HASH unknown
24+ #endif
25+
26+ #define STR_HELPER (x ) #x
27+ #define STR (x ) STR_HELPER (x)
28+
29+
1930MLLM_MAIN ({
2031 auto & help = mllm::Argparse::add<bool >(" -h|--help" ).help (" Show help message" );
2132 auto & model_name = mllm::Argparse::add<std::string>(" -n|--model_name" ).help (" Model name" );
@@ -25,8 +36,19 @@ MLLM_MAIN({
2536 auto & pp = mllm::Argparse::add<std::string>(" -pp|--prompt_length" ).help (" Prompt length" );
2637 auto & tg = mllm::Argparse::add<std::string>(" -tg|--test_generation_length" ).help (" Test Generation length" );
2738 auto & cache_length = mllm::Argparse::add<int32_t >(" -cl|--cache_length" ).help (" Cache length" );
39+
40+ // New CLI Arguments
41+ auto & runs = mllm::Argparse::add<int32_t >(" -r|--runs" ).help (" Number of benchmark runs" ).def (3 );
42+ auto & cooldown_s = mllm::Argparse::add<int32_t >(" -cs|--cooldown_s" ).help (" Cooldown time between runs in seconds" ).def (5 );
43+ auto & output_csv = mllm::Argparse::add<std::string>(" -oc|--output_csv" ).help (" Output results to a CSV file" ).def (" " );
44+ auto & schema_version = mllm::Argparse::add<int32_t >(" -sv|--schema_version" ).help (" Schema version for output format" ).def (1 );
45+ auto & kv_dtype_bytes = mllm::Argparse::add<int32_t >(" -kv|--kv_dtype_bytes" ).help (" KV cache data type bytes (1: int8, 2: fp16, 4: fp32)" ).def (4 );
46+
2847 mllm::Argparse::parse (argc, argv);
2948
49+ mllm::Context::instance ().setCpuOpThreads (num_threads.get ());
50+ mllm::setMaximumNumThreads ((uint32_t )num_threads.get ());
51+
3052 // Print Build Version
3153 mllm::print (" MLLM Build Version :" , STRINGIFY (MLLM_GIT_COMMIT_HASH ));
3254
@@ -58,6 +80,25 @@ MLLM_MAIN({
5880 auto benchmark = createBenchmark (model_name.get ());
5981 MLLM_RT_ASSERT (benchmark != nullptr );
6082
83+
84+ // Validate runs early to avoid huge reserve() when negative values cast to size_t.
85+ int R = runs.get ();
86+ if (R <= 0 ) {
87+ mllm::print (" [ERROR] --runs must be > 0, got:" , R);
88+ return 1 ;
89+ }
90+
91+ // Open file stream
92+ std::ofstream csv_file;
93+ if (!output_csv.get ().empty ()) {
94+ csv_file.open (output_csv.get ());
95+ if (!csv_file.is_open ()) {
96+ mllm::print (" [ERROR] Failed to open --output_csv:" , output_csv.get ());
97+ return 1 ;
98+ }
99+ csv_file << " schema_version,git_commit,arch,model_name,pp,tg,ttft_ms,prefill_speed,decode_speed,prefill_ms,decode_ms_per_tok,kv_est_bytes_pp,kv_est_bytes_final\n " ;
100+ }
101+
61102 // Print Model Info
62103 mllm::print (" Model Info" );
63104 benchmark->init (config_path.get (), model_path.get (), cache_length.get ());
@@ -92,7 +133,7 @@ MLLM_MAIN({
92133 for (size_t i = 0 ; i < pp_values.size (); ++i) { pp_tg_pairs.emplace_back (pp_values[i], tg_values[i]); }
93134 }
94135
95- // Actual run for 3 turns and gives avg results. Each turn will sleep for 5 seconds to let the SoC or GPU/NPU cool down.
136+ // Actual run for configurable number of turns
96137 mllm::print (" \n ========================================" );
97138 mllm::print (" Starting Benchmark Tests" );
98139 mllm::print (" ========================================\n " );
@@ -106,30 +147,40 @@ MLLM_MAIN({
106147
107148 // Storage for results
108149 std::vector<BenchmarkTemplateResult> results;
109- results.reserve (3 );
150+ results.reserve (static_cast < size_t >(R) );
110151
111- for (int i = 0 ; i < 3 ; ++i) {
112- mllm::print (" Run" , i + 1 , " of 3 ..." );
152+ for (int i = 0 ; i < R ; ++i) {
153+ mllm::print (" Run" , i + 1 , " of" , R, " ..." );
113154
114- // Clear cache before each run
115- benchmark->clear ();
155+ // Clear cache/state before each run to reduce cross-run interference.
116156
117- // Run benchmark
157+ benchmark->clear ();
158+ // Run benchmark for this (pp, tg) pair.
118159 auto result = benchmark->run (pp, tg);
119160 results.push_back (result);
120161
121162 mllm::print (" TTFT :" , result.ttft , " ms" );
122163 mllm::print (" Prefill Speed:" , result.prefill_speed , " tokens/s" );
123164 mllm::print (" Decode Speed :" , result.decode_speed , " tokens/s" );
124165
125- // Sleep for 5 seconds between runs to cool down
126- if (i < 2 ) {
127- mllm::print (" Cooling down for 5 seconds..." );
128- std::this_thread::sleep_for (std::chrono::seconds (5 ));
166+ // Derive per-run latency numbers from throughput (guard against divide-by-zero).
167+
168+ float prefill_ms = (result.prefill_speed > 0 .0f ) ? (pp / result.prefill_speed ) * 1000 .0f : 0 .0f ;
169+ float decode_ms_per_tok = (result.decode_speed > 0 .0f ) ? (1 .0f / result.decode_speed ) * 1000 .0f : 0 .0f ;
170+ mllm::print (" Prefill Latency :" , prefill_ms, " ms" );
171+ mllm::print (" Decode Latency :" , decode_ms_per_tok, " ms" );
172+
173+ // Sleep between runs to cool down (configurable).
174+
175+ int cool = cooldown_s.get ();
176+ if (i + 1 < R && cool > 0 ) {
177+ mllm::print (" Cooling down for" , cool, " seconds..." );
178+ std::this_thread::sleep_for (std::chrono::seconds (cool));
129179 }
130180 }
131181
132182 // Calculate average results
183+ float denom = (R > 0 ) ? static_cast <float >(R) : 1 .0f ;
133184 float avg_ttft = 0 .0f ;
134185 float avg_prefill_speed = 0 .0f ;
135186 float avg_decode_speed = 0 .0f ;
@@ -140,20 +191,44 @@ MLLM_MAIN({
140191 avg_decode_speed += result.decode_speed ;
141192 }
142193
143- avg_ttft /= 3 .0f ;
144- avg_prefill_speed /= 3 .0f ;
145- avg_decode_speed /= 3 .0f ;
146-
147- // Print average results
148- mllm::print (" \n ========== Average Results ==========" );
149- mllm::print (" Configuration: PP=" , pp, " TG=" , tg);
150- mllm::print (" Average TTFT :" , avg_ttft, " ms" );
151- mllm::print (" Average Prefill Speed:" , avg_prefill_speed, " tokens/s" );
152- mllm::print (" Average Decode Speed :" , avg_decode_speed, " tokens/s" );
153- mllm::print (" =====================================\n " );
194+ avg_ttft /= denom;
195+ avg_prefill_speed /= denom;
196+ avg_decode_speed /= denom;
197+
198+ float avg_prefill_ms = (avg_prefill_speed > 0 .0f ) ? (pp / avg_prefill_speed) * 1000 .0f : 0 .0f ;
199+ float avg_decode_ms_per_tok = (avg_decode_speed > 0 .0f ) ? (1 .0f / avg_decode_speed) * 1000 .0f : 0 .0f ;
200+
201+ // Rough KV cache estimate (bytes)
202+ double kv_est_bytes_pp = 0.0 ;
203+ double kv_est_bytes_final = 0.0 ;
204+
205+ // Prepare one line output (avg)
206+ std::stringstream ss;
207+ ss << schema_version.get () << " ,"
208+ << STRINGIFY (MLLM_GIT_COMMIT_HASH ) << " ,"
209+ << mllm::cpu::CURRENT_ARCH_STRING << " ,"
210+ << model_name.get () << " ,"
211+ << pp << " ,"
212+ << tg << " ,"
213+ << avg_ttft << " ,"
214+ << avg_prefill_speed << " ,"
215+ << avg_decode_speed << " ,"
216+ << avg_prefill_ms << " ,"
217+ << avg_decode_ms_per_tok << " ,"
218+ << kv_est_bytes_pp << " ,"
219+ << kv_est_bytes_final;
220+
221+ if (csv_file.is_open ()) {
222+ csv_file << ss.str () << std::endl;
223+ }
154224 }
155225
156226 mllm::print (" \n ========================================" );
157227 mllm::print (" Benchmark Tests Completed" );
158228 mllm::print (" ========================================" );
229+
230+ // close file stream
231+ if (csv_file.is_open ()) {
232+ csv_file.close ();
233+ }
159234})
0 commit comments