@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a0b782f9cc
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 11, 06:15 UTC

[Metrics] Model FLOPs Utilization estimation (#30738)

Signed-off-by: SungMinCho <tjdals4565@gmail.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com>

SungMinCho committed Dec 17, 2025 at 17:40 UTC a0b782f9ccd02add3516f074de163c9847686066
8 files changed +2186 -2
tests/v1/metrics/test_perf_metrics.py new
+897
@@ -0,0 +1,897 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Tests for the analytic estimators in metrics/flops.py.
5 +"""
6 +
7 +import types
8 +from types import SimpleNamespace
9 +
10 +from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config
11 +from transformers.models.llama4.configuration_llama4 import (
12 + Llama4Config,
13 + Llama4TextConfig,
14 +)
15 +from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
16 +from transformers.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig
17 +
18 +from vllm.config.model import ModelConfig, get_hf_text_config
19 +from vllm.v1.metrics.perf import (
20 + AttentionMetrics,
21 + BaseConfigParser,
22 + ExecutionContext,
23 + FfnMetrics,
24 + ModelMetrics,
25 + ParsedArgs,
26 + UnembedMetrics,
27 +)
28 +
29 +
30 +class MockModelConfig:
31 + """Mock ModelConfig that implements the getter methods used by parsers."""
32 +
33 + def __init__(self, hf_config, dtype):
34 + self.hf_config = hf_config
35 + self.hf_text_config = get_hf_text_config(hf_config)
36 + self.dtype = dtype
37 + self.is_attention_free = False
38 +
39 + def __getattr__(self, name):
40 + # 1. Check if ModelConfig actually has this attribute
41 + if not hasattr(ModelConfig, name):
42 + raise AttributeError(
43 + f"'{type(self).__name__}' object has no attribute '{name}' "
44 + f"and neither does 'ModelConfig'."
45 + )
46 +
47 + # 2. Fetch the attribute from the ModelConfig CLASS
48 + attr = getattr(ModelConfig, name)
49 +
50 + # 3. Case A: It is a @property
51 + if isinstance(attr, property):
52 + # Manually invoke the property's getter, passing 'self' (this mock instance)
53 + return attr.__get__(self, self.__class__)
54 +
55 + # 4. Case B: It is a standard method (function)
56 + if isinstance(attr, types.FunctionType):
57 + # Bind the function to 'self' so it acts like a method of
58 + # this instance. This creates a bound method where 'self' is
59 + # automatically passed as the first arg.
60 + return types.MethodType(attr, self)
61 +
62 + # 5. Case C: It is a class attribute / static variable
63 + return attr
64 +
65 +
66 +def create_mock_vllm_config(
67 + hf_config,
68 + model_dtype="bfloat16",
69 + cache_dtype="auto",
70 + quant_config=None,
71 + data_parallel_size=1,
72 + tensor_parallel_size=1,
73 + pipeline_parallel_size=1,
74 + enable_expert_parallel=False,
75 +) -> SimpleNamespace:
76 + vllm_config = SimpleNamespace()
77 + vllm_config.model_config = MockModelConfig(hf_config, model_dtype)
78 +
79 + vllm_config.cache_config = SimpleNamespace()
80 + vllm_config.cache_config.cache_dtype = cache_dtype
81 +
82 + vllm_config.quant_config = quant_config
83 +
84 + vllm_config.parallel_config = SimpleNamespace()
85 + vllm_config.parallel_config.data_parallel_size = data_parallel_size
86 + vllm_config.parallel_config.tensor_parallel_size = tensor_parallel_size
87 + vllm_config.parallel_config.pipeline_parallel_size = pipeline_parallel_size
88 + vllm_config.parallel_config.enable_expert_parallel = enable_expert_parallel
89 +
90 + return vllm_config
91 +
92 +
93 +#### Parser Tests ####
94 +
95 +
96 +def test_base_config_parser():
97 + """Test BaseConfigParser extracts base model attributes correctly."""
98 + hf_config = Qwen3Config(
99 + vocab_size=50000,
100 + hidden_size=2048,
101 + num_attention_heads=16,
102 + num_hidden_layers=24,
103 + )
104 + vllm_config = create_mock_vllm_config(hf_config, model_dtype="float16")
105 +
106 + parser = BaseConfigParser()
107 + args = ParsedArgs()
108 + result = parser.parse(args, vllm_config)
109 +
110 + assert result.vocab_size == 50000
111 + assert result.hidden_size == 2048
112 + assert result.num_attention_heads == 16
113 + assert result.num_hidden_layers == 24
114 + assert result.weight_byte_size == 2 # float16 is 2 bytes
115 + assert result.activation_byte_size == 2 # default activation size
116 +
117 +
118 +def test_base_attention_config_parser_with_gqa():
119 + """Test BaseAttentionConfigParser with grouped query attention."""
120 + hf_config = Qwen3Config(
121 + hidden_size=4096,
122 + num_attention_heads=32,
123 + num_key_value_heads=8, # GQA with 4:1 ratio
124 + head_dim=128,
125 + )
126 + vllm_config = create_mock_vllm_config(hf_config)
127 +
128 + parser_chain = AttentionMetrics.get_parser()
129 + result = parser_chain.parse(vllm_config)
130 +
131 + assert result.num_key_value_heads == 8
132 + assert result.head_dim == 128
133 +
134 +
135 +def test_base_attention_config_parser_without_gqa():
136 + """
137 + Test BaseAttentionConfigParser defaults to MHA when num_key_value_heads not
138 + specified.
139 + """
140 + hf_config = Qwen3Config(
141 + hidden_size=4096,
142 + num_attention_heads=32,
143 + # No num_key_value_heads specified
144 + )
145 + vllm_config = create_mock_vllm_config(hf_config)
146 +
147 + parser_chain = AttentionMetrics.get_parser()
148 + result = parser_chain.parse(vllm_config)
149 +
150 + # Should default to MHA (num_key_value_heads = num_attention_heads)
151 + assert result.num_key_value_heads == 32
152 +
153 +
154 +def test_base_ffn_config_parser_dense():
155 + """Test BaseFfnConfigParser for dense FFN."""
156 + hf_config = Qwen3Config(
157 + hidden_size=4096,
158 + intermediate_size=11008,
159 + num_hidden_layers=32,
160 + )
161 + vllm_config = create_mock_vllm_config(hf_config)
162 +
163 + parser_chain = FfnMetrics.get_parser()
164 + result = parser_chain.parse(vllm_config)
165 +
166 + assert result.intermediate_size == 11008
167 + assert result.num_experts == 0
168 + assert result.num_experts_per_tok == 0
169 + assert result.num_moe_layers == 0 # No MoE
170 +
171 +
172 +def test_base_ffn_config_parser_moe():
173 + """Test BaseFfnConfigParser for MoE FFN."""
174 + hf_config = Qwen3MoeConfig(
175 + hidden_size=4096,
176 + intermediate_size=11008,
177 + num_hidden_layers=32,
178 + num_experts=64,
179 + num_experts_per_tok=8,
180 + moe_intermediate_size=14336,
181 + n_shared_experts=2,
182 + )
183 + vllm_config = create_mock_vllm_config(hf_config)
184 +
185 + parser_chain = FfnMetrics.get_parser()
186 + result = parser_chain.parse(vllm_config)
187 +
188 + assert result.num_experts == 64
189 + assert result.num_experts_per_tok == 8
190 + assert result.moe_intermediate_size == 14336
191 + assert result.num_shared_experts == 2
192 + assert result.num_moe_layers == 32 # All layers are MoE by default
193 +
194 +
195 +def test_interleave_moe_layer_step_parser():
196 + """Test InterleaveMoeLayerStepParser correctly computes MoE layer count."""
197 + hf_config = Llama4Config(
198 + text_config=Llama4TextConfig(
199 + num_hidden_layers=32,
200 + num_local_experts=64,
201 + interleave_moe_layer_step=4, # Every 4th layer is MoE
202 + ),
203 + )
204 +
205 + vllm_config = create_mock_vllm_config(hf_config)
206 +
207 + parser_chain = FfnMetrics.get_parser()
208 + result = parser_chain.parse(vllm_config)
209 +
210 + assert result.num_moe_layers == 8
211 +
212 +
213 +def test_moe_layer_freq_parser():
214 + """Test MoeLayerFreqParser correctly computes MoE layer count."""
215 + hf_config = DeepseekV3Config(
216 + num_hidden_layers=30,
217 + n_routed_experts=64,
218 + moe_layer_freq=3, # Every 3rd layer after first_k_dense_replace
219 + first_k_dense_replace=6, # First 6 layers are dense
220 + )
221 + vllm_config = create_mock_vllm_config(hf_config)
222 +
223 + parser_chain = FfnMetrics.get_parser()
224 + result = parser_chain.parse(vllm_config)
225 +
226 + # Layers >= 6 and divisible by 3: 6, 9, 12, 15, 18, 21, 24, 27
227 + expected_moe_layers = len(
228 + [layer for layer in range(30) if layer >= 6 and layer % 3 == 0]
229 + )
230 + assert expected_moe_layers == 8
231 + assert result.num_moe_layers == expected_moe_layers
232 +
233 +
234 +#### ComponentMetrics Tests ####
235 +
236 +
237 +def test_attention_metrics_scaling():
238 + """Test that attention metrics scale proportionally with model dimensions."""
239 + base_hf_config = Qwen3Config(
240 + hidden_size=2048,
241 + num_attention_heads=16,
242 + num_key_value_heads=16,
243 + num_hidden_layers=12,
244 + head_dim=128,
245 + )
246 +
247 + base_vllm_config = create_mock_vllm_config(base_hf_config)
248 + base_metrics = AttentionMetrics.from_vllm_config(base_vllm_config)
249 +
250 + # Test scaling with number of layers
251 + double_layers_hf_config = Qwen3Config(
252 + hidden_size=2048,
253 + num_attention_heads=16,
254 + num_key_value_heads=16,
255 + num_hidden_layers=24, # Double the layers
256 + head_dim=128,
257 + )
258 + double_layers_vllm_config = create_mock_vllm_config(double_layers_hf_config)
259 + double_layers_metrics = AttentionMetrics.from_vllm_config(double_layers_vllm_config)
260 +
261 + ctx = ExecutionContext.from_single_request(
262 + num_tokens=100, context_len=512, is_prefill=True
263 + )
264 +
265 + # FLOPS should double when layers double
266 + base_flops = base_metrics.get_num_flops(ctx)
267 + double_flops = double_layers_metrics.get_num_flops(ctx)
268 + assert double_flops == 2 * base_flops
269 +
270 + # Read/write bytes should also scale proportionally
271 + base_read = base_metrics.get_read_bytes(ctx)
272 + double_read = double_layers_metrics.get_read_bytes(ctx)
273 + assert double_read == 2 * base_read
274 +
275 + base_write = base_metrics.get_write_bytes(ctx)
276 + double_write = double_layers_metrics.get_write_bytes(ctx)
277 + assert double_write == 2 * base_write
278 +
279 +
280 +def test_attention_metrics_grouped_query():
281 + """Test attention metrics handle grouped query attention correctly."""
282 + mha_hf_config = Qwen3Config(
283 + hidden_size=4096,
284 + num_attention_heads=32,
285 + num_key_value_heads=32, # MHA
286 + num_hidden_layers=1,
287 + )
288 + mha_config = create_mock_vllm_config(mha_hf_config)
289 +
290 + gqa_hf_config = Qwen3Config(
291 + hidden_size=4096,
292 + num_attention_heads=32,
293 + num_key_value_heads=8, # GQA with 4:1 ratio
294 + num_hidden_layers=1,
295 + )
296 + gqa_config = create_mock_vllm_config(gqa_hf_config)
297 +
298 + mha_metrics = AttentionMetrics.from_vllm_config(mha_config)
299 + gqa_metrics = AttentionMetrics.from_vllm_config(gqa_config)
300 +
301 + ctx = ExecutionContext.from_single_request(
302 + num_tokens=1, context_len=1024, is_prefill=False
303 + )
304 +
305 + # GQA should have less KV cache reads since fewer KV heads
306 + mha_read = mha_metrics.get_read_bytes(ctx)
307 + gqa_read = gqa_metrics.get_read_bytes(ctx)
308 + assert gqa_read < mha_read
309 +
310 +
311 +def test_ffn_metrics_scaling():
312 + """Test FFN metrics scale proportionally with model dimensions."""
313 + base_hf_config = Qwen3Config(
314 + hidden_size=2048,
315 + intermediate_size=8192,
316 + num_hidden_layers=12,
317 + )
318 + base_vllm_config = create_mock_vllm_config(base_hf_config)
319 + base_metrics = FfnMetrics.from_vllm_config(base_vllm_config)
320 +
321 + # Test scaling with intermediate size
322 + larger_ffn_hf_config = Qwen3Config(
323 + hidden_size=2048,
324 + intermediate_size=16384, # Double intermediate size
325 + num_hidden_layers=12,
326 + )
327 + larger_ffn_vllm_config = create_mock_vllm_config(larger_ffn_hf_config)
328 + larger_ffn_metrics = FfnMetrics.from_vllm_config(larger_ffn_vllm_config)
329 +
330 + ctx = ExecutionContext.from_single_request(
331 + num_tokens=100, context_len=512, is_prefill=True
332 + )
333 +
334 + # FLOPS should double when intermediate size doubles
335 + base_flops = base_metrics.get_num_flops(ctx)
336 + larger_flops = larger_ffn_metrics.get_num_flops(ctx)
337 + assert larger_flops == base_flops * 2
338 +
339 +
340 +def test_moe_metrics_vs_dense():
341 + """Test MoE metrics versus dense metrics."""
342 + dense_hf_config = Qwen3Config(
343 + hidden_size=2048,
344 + intermediate_size=8192,
345 + num_hidden_layers=12,
346 + )
347 + dense_config = create_mock_vllm_config(dense_hf_config)
348 +
349 + moe_hf_config = Qwen3MoeConfig(
350 + hidden_size=2048,
351 + intermediate_size=8192,
352 + num_hidden_layers=12,
353 + num_experts=64,
354 + num_experts_per_tok=2, # 2 routed expert
355 + moe_intermediate_size=8192,
356 + n_shared_experts=0,
357 + )
358 + moe_config = create_mock_vllm_config(moe_hf_config)
359 +
360 + dense_metrics = FfnMetrics.from_vllm_config(dense_config)
361 + moe_metrics = FfnMetrics.from_vllm_config(moe_config)
362 +
363 + ctx = ExecutionContext.from_single_request(
364 + num_tokens=100, context_len=512, is_prefill=True
365 + )
366 +
367 + # MoE should have different compute/memory characteristics
368 + dense_flops = dense_metrics.get_num_flops(ctx)
369 + moe_flops = moe_metrics.get_num_flops(ctx)
370 +
371 + # 2 routed experts vs 1 dense.
372 + assert moe_flops == dense_flops * 2
373 +
374 +
375 +def test_unembed_metrics_scaling():
376 + """Test unembedding metrics scale with vocab size."""
377 + small_vocab_hf_config = Qwen3Config(
378 + hidden_size=2048,
379 + vocab_size=32000,
380 + )
381 + small_vocab_config = create_mock_vllm_config(small_vocab_hf_config)
382 +
383 + large_vocab_hf_config = Qwen3Config(
384 + hidden_size=2048,
385 + vocab_size=64000, # Double vocab size
386 + )
387 + large_vocab_config = create_mock_vllm_config(large_vocab_hf_config)
388 +
389 + small_vocab_metrics = UnembedMetrics.from_vllm_config(small_vocab_config)
390 + large_vocab_metrics = UnembedMetrics.from_vllm_config(large_vocab_config)
391 +
392 + ctx = ExecutionContext.from_single_request(
393 + num_tokens=100, context_len=512, is_prefill=True
394 + )
395 +
396 + # FLOPS should double when vocab size doubles
397 + small_flops = small_vocab_metrics.get_num_flops(ctx)
398 + large_flops = large_vocab_metrics.get_num_flops(ctx)
399 + assert large_flops == 2 * small_flops
400 +
401 +
402 +def test_prefill_vs_decode_differences():
403 + """Test that prefill and decode have different memory access patterns."""
404 + hf_config = Qwen3Config(
405 + hidden_size=2048,
406 + num_attention_heads=16,
407 + num_key_value_heads=16,
408 + num_hidden_layers=1,
409 + )
410 + config = create_mock_vllm_config(hf_config)
411 +
412 + metrics = AttentionMetrics.from_vllm_config(config)
413 +
414 + prefill_ctx = ExecutionContext.from_single_request(
415 + num_tokens=512, context_len=512, is_prefill=True
416 + )
417 + decode_ctx = ExecutionContext.from_single_request(
418 + num_tokens=1, context_len=512, is_prefill=False
419 + )
420 +
421 + prefill_read = metrics.get_read_bytes(prefill_ctx)
422 + decode_read = metrics.get_read_bytes(decode_ctx)
423 +
424 + assert prefill_read != decode_read
425 +
426 +
427 +def test_model_metrics_aggregation():
428 + """Test ModelMetrics correctly aggregates across components."""
429 + hf_config = Qwen3Config(
430 + hidden_size=2048,
431 + num_attention_heads=16,
432 + num_hidden_layers=12,
433 + vocab_size=32000,
434 + intermediate_size=8192,
435 + )
436 + config = create_mock_vllm_config(hf_config)
437 +
438 + model_metrics = ModelMetrics(config)
439 + ctx = ExecutionContext.from_single_request(
440 + num_tokens=100, context_len=512, is_prefill=True
441 + )
442 +
443 + # Should have metrics for attention, ffn, and unembed
444 + total_flops = model_metrics.get_num_flops(ctx)
445 + breakdown = model_metrics.get_num_flops_breakdown(ctx)
446 +
447 + # Breakdown should sum to total
448 + assert total_flops == sum(breakdown.values())
449 +
450 +
451 +def test_moe_expert_activation_proportional_scaling():
452 + """Test that routed expert metrics scale proportionally with num_experts_per_tok."""
453 + base_moe_config = Qwen3MoeConfig(
454 + hidden_size=2048,
455 + intermediate_size=8192,
456 + num_hidden_layers=12,
457 + num_experts=64,
458 + num_experts_per_tok=1, # 1 expert per token
459 + moe_intermediate_size=8192,
460 + n_shared_experts=2,
461 + )
462 +
463 + double_experts_config = Qwen3MoeConfig(
464 + hidden_size=2048,
465 + intermediate_size=8192,
466 + num_hidden_layers=12,
467 + num_experts=64,
468 + num_experts_per_tok=2, # 2 experts per token (double)
469 + moe_intermediate_size=8192,
470 + n_shared_experts=2, # Same shared experts
471 + )
472 +
473 + triple_experts_config = Qwen3MoeConfig(
474 + hidden_size=2048,
475 + intermediate_size=8192,
476 + num_hidden_layers=12,
477 + num_experts=64,
478 + num_experts_per_tok=3, # 3 experts per token (triple)
479 + moe_intermediate_size=8192,
480 + n_shared_experts=2, # Same shared experts
481 + )
482 +
483 + base_vllm_config = create_mock_vllm_config(base_moe_config)
484 + double_vllm_config = create_mock_vllm_config(double_experts_config)
485 + triple_vllm_config = create_mock_vllm_config(triple_experts_config)
486 +
487 + base_metrics = FfnMetrics.from_vllm_config(base_vllm_config)
488 + double_metrics = FfnMetrics.from_vllm_config(double_vllm_config)
489 + triple_metrics = FfnMetrics.from_vllm_config(triple_vllm_config)
490 +
491 + ctx = ExecutionContext.from_single_request(
492 + num_tokens=100, context_len=512, is_prefill=True
493 + )
494 +
495 + # Get total metrics - the key insight is that differences should be proportional
496 + base_flops = base_metrics.get_num_flops(ctx)
497 + double_flops = double_metrics.get_num_flops(ctx)
498 + triple_flops = triple_metrics.get_num_flops(ctx)
499 +
500 + # The difference between double and base should equal one additional expert
501 + one_expert_diff = double_flops - base_flops
502 +
503 + # The difference between triple and base should equal two additional experts
504 + two_expert_diff = triple_flops - base_flops
505 +
506 + # Proportional scaling: 2 * (1 expert diff) should equal (2 expert diff)
507 + assert two_expert_diff == 2 * one_expert_diff
508 +
509 + # Same logic applies to memory operations
510 + base_read = base_metrics.get_read_bytes(ctx)
511 + double_read = double_metrics.get_read_bytes(ctx)
512 + triple_read = triple_metrics.get_read_bytes(ctx)
513 +
514 + one_expert_read_diff = double_read - base_read
515 + two_expert_read_diff = triple_read - base_read
516 +
517 + assert two_expert_read_diff == 2 * one_expert_read_diff
518 +
519 + # Same for write bytes
520 + base_write = base_metrics.get_write_bytes(ctx)
521 + double_write = double_metrics.get_write_bytes(ctx)
522 + triple_write = triple_metrics.get_write_bytes(ctx)
523 +
524 + one_expert_write_diff = double_write - base_write
525 + two_expert_write_diff = triple_write - base_write
526 +
527 + assert two_expert_write_diff == 2 * one_expert_write_diff
528 +
529 +
530 +def test_quantization_config_parser_fp8():
531 + """Test quantization parsers with fp8."""
532 +
533 + class MockQuantConfig:
534 + def get_name(self):
535 + return "fp8"
536 +
537 + hf_config = Qwen3Config(
538 + hidden_size=2048, num_attention_heads=16, num_hidden_layers=1
539 + )
540 + vllm_config = create_mock_vllm_config(hf_config, quant_config=MockQuantConfig())
541 +
542 + attn_result = AttentionMetrics.get_parser().parse(vllm_config)
543 + assert attn_result.weight_byte_size == 1 # fp8
544 +
545 + ffn_result = FfnMetrics.get_parser().parse(vllm_config)
546 + assert ffn_result.weight_byte_size == 1 # fp8
547 +
548 +
549 +def test_quantization_config_parser_mxfp4():
550 + """Test quantization parsers with mxfp4."""
551 +
552 + class MockQuantConfig:
553 + def get_name(self):
554 + return "mxfp4"
555 +
556 + hf_config = Qwen3Config(
557 + hidden_size=2048, intermediate_size=8192, num_hidden_layers=1
558 + )
559 + vllm_config = create_mock_vllm_config(hf_config, quant_config=MockQuantConfig())
560 +
561 + ffn_result = FfnMetrics.get_parser().parse(vllm_config)
562 + assert ffn_result.weight_byte_size == 0.5 # mxfp4
563 +
564 +
565 +#### Per-GPU Tests ####
566 +
567 +
568 +def test_attention_per_gpu_with_tensor_parallelism():
569 + """Test attention metrics with tensor parallelism - per_gpu vs global."""
570 + hf_config = Qwen3Config(
571 + hidden_size=4096,
572 + num_attention_heads=32,
573 + num_key_value_heads=8,
574 + num_hidden_layers=24,
575 + )
576 +
577 + # Test with TP=4
578 + vllm_config = create_mock_vllm_config(hf_config, tensor_parallel_size=4)
579 + metrics = AttentionMetrics.from_vllm_config(vllm_config)
580 +
581 + ctx = ExecutionContext.from_single_request(
582 + num_tokens=128, context_len=1024, is_prefill=True
583 + )
584 +
585 + # Get global and per-gpu metrics
586 + global_flops = metrics.get_num_flops(ctx, per_gpu=False)
587 + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True)
588 +
589 + # With TP=4, global flops should be 4x per-gpu flops (heads divided by 4)
590 + assert global_flops == 4 * per_gpu_flops
591 +
592 + # Same for read/write bytes
593 + global_read = metrics.get_read_bytes(ctx, per_gpu=False)
594 + per_gpu_read = metrics.get_read_bytes(ctx, per_gpu=True)
595 + # Reads should scale similarly (weight reads are divided by TP)
596 + assert global_read > per_gpu_read
597 +
598 + global_write = metrics.get_write_bytes(ctx, per_gpu=False)
599 + per_gpu_write = metrics.get_write_bytes(ctx, per_gpu=True)
600 + assert global_write > per_gpu_write
601 +
602 +
603 +def test_attention_per_gpu_with_pipeline_parallelism():
604 + """Test attention metrics with pipeline parallelism - per_gpu vs global."""
605 + hf_config = Qwen3Config(
606 + hidden_size=2048,
607 + num_attention_heads=16,
608 + num_hidden_layers=32,
609 + )
610 +
611 + # Test with PP=4
612 + vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=4)
613 + metrics = AttentionMetrics.from_vllm_config(vllm_config)
614 +
615 + ctx = ExecutionContext.from_single_request(
616 + num_tokens=100, context_len=512, is_prefill=False
617 + )
618 +
619 + # Get global and per-gpu metrics
620 + global_flops = metrics.get_num_flops(ctx, per_gpu=False)
621 + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True)
622 +
623 + # With PP=4, global flops should be 4x per-gpu flops (layers divided by 4)
624 + assert global_flops == 4 * per_gpu_flops
625 +
626 + global_read = metrics.get_read_bytes(ctx, per_gpu=False)
627 + per_gpu_read = metrics.get_read_bytes(ctx, per_gpu=True)
628 + assert global_read == 4 * per_gpu_read
629 +
630 +
631 +def test_ffn_per_gpu_with_tensor_parallelism():
632 + """Test FFN metrics with tensor parallelism - per_gpu vs global."""
633 + hf_config = Qwen3Config(
634 + hidden_size=4096,
635 + intermediate_size=14336,
636 + num_hidden_layers=32,
637 + )
638 +
639 + # Test with DP=2, TP=4 (ffn_tp_size will be 8)
640 + vllm_config = create_mock_vllm_config(
641 + hf_config,
642 + data_parallel_size=2,
643 + tensor_parallel_size=4,
644 + )
645 + metrics = FfnMetrics.from_vllm_config(vllm_config)
646 +
647 + # ffn_tp_size should be dp_size * tp_size = 8 (when EP not enabled)
648 + assert metrics.ffn_tp_size == 8
649 +
650 + ctx = ExecutionContext.from_single_request(
651 + num_tokens=128, context_len=2048, is_prefill=True
652 + )
653 +
654 + # Get global and per-gpu metrics
655 + global_flops = metrics.get_num_flops(ctx, per_gpu=False)
656 + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True)
657 +
658 + # With ffn_tp_size=8, global should be 8x per-gpu
659 + assert global_flops == 8 * per_gpu_flops
660 +
661 +
662 +def test_ffn_per_gpu_with_pipeline_parallelism():
663 + """Test FFN metrics with pipeline parallelism - per_gpu vs global."""
664 + hf_config = Qwen3Config(
665 + hidden_size=2048,
666 + intermediate_size=8192,
667 + num_hidden_layers=24,
668 + )
669 +
670 + # Test with PP=6
671 + vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=6)
672 + metrics = FfnMetrics.from_vllm_config(vllm_config)
673 +
674 + ctx = ExecutionContext.from_single_request(
675 + num_tokens=100, context_len=512, is_prefill=True
676 + )
677 +
678 + # Get global and per-gpu metrics
679 + global_flops = metrics.get_num_flops(ctx, per_gpu=False)
680 + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True)
681 +
682 + # With PP=6, global should be 6x per-gpu (layers divided by 6)
683 + assert global_flops == 6 * per_gpu_flops
684 +
685 +
686 +def test_moe_per_gpu_with_expert_parallelism():
687 + """
688 + Test MoE metrics with expert parallelism - verifies num_activated_experts bug fix.
689 + """
690 + hf_config = Qwen3MoeConfig(
691 + hidden_size=2048,
692 + intermediate_size=8192,
693 + num_hidden_layers=24,
694 + num_experts=64,
695 + num_experts_per_tok=8,
696 + moe_intermediate_size=14336,
697 + n_shared_experts=2,
698 + )
699 +
700 + # Test with DP=2, TP=4, EP enabled (ffn_ep_size will be 8)
701 + vllm_config = create_mock_vllm_config(
702 + hf_config,
703 + data_parallel_size=2,
704 + tensor_parallel_size=4,
705 + enable_expert_parallel=True,
706 + )
707 + metrics = FfnMetrics.from_vllm_config(vllm_config)
708 +
709 + # When EP enabled, ffn_ep_size = dp_size * tp_size = 8
710 + assert metrics.ffn_ep_size == 8
711 + assert metrics.ffn_tp_size == 1
712 +
713 + ctx = ExecutionContext.from_single_request(
714 + num_tokens=100, context_len=512, is_prefill=True
715 + )
716 +
717 + # Get per-gpu metrics
718 + per_gpu_read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=True)
719 + global_read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=False)
720 +
721 + # Verify that routed expert weight reads are reasonable
722 + # With per_gpu=True, each GPU has 64/8 = 8 experts
723 + # T=100, E_per_gpu=8/8=1, so T*E=100 expert activations
724 + # num_activated_experts should be min(100, 8) = 8
725 +
726 + # Check that weight reads scale appropriately
727 + # Global has all 64 experts, per-gpu has 8 experts
728 + # So weight reads should reflect this difference
729 + if "routed_up_gate_weights" in per_gpu_read_breakdown:
730 + per_gpu_weight_reads = per_gpu_read_breakdown["routed_up_gate_weights"]
731 + global_weight_reads = global_read_breakdown["routed_up_gate_weights"]
732 +
733 + # The ratio should reflect the expert count difference
734 + # This verifies the bug fix works correctly
735 + assert per_gpu_weight_reads < global_weight_reads
736 +
737 + # Global should read more experts than per-gpu
738 + # Exact ratio depends on num_activated_experts calculation
739 + ratio = global_weight_reads / per_gpu_weight_reads
740 + # Should be > 1 since global has more experts to read
741 + assert ratio > 1
742 +
743 +
744 +def test_moe_per_gpu_expert_activation_accounting():
745 + """
746 + Test that MoE correctly accounts for expert activations with small batch sizes.
747 + """
748 + hf_config = Qwen3MoeConfig(
749 + hidden_size=2048,
750 + intermediate_size=8192,
751 + num_hidden_layers=12,
752 + num_experts=64,
753 + num_experts_per_tok=8,
754 + moe_intermediate_size=14336,
755 + n_shared_experts=0, # No shared experts for this test
756 + )
757 +
758 + # Test with EP=8
759 + vllm_config = create_mock_vllm_config(
760 + hf_config,
761 + data_parallel_size=8,
762 + enable_expert_parallel=True,
763 + )
764 + metrics = FfnMetrics.from_vllm_config(vllm_config)
765 +
766 + # Small batch: T=10, E_per_gpu=8/8=1
767 + # Each GPU: T*E = 10*1 = 10 activations
768 + # Experts per GPU: 64/8 = 8
769 + # So num_activated_experts should be min(10, 8) = 8
770 + small_ctx = ExecutionContext.from_single_request(
771 + num_tokens=10, context_len=512, is_prefill=True
772 + )
773 + small_read = metrics.get_read_bytes_breakdown(small_ctx, per_gpu=True)
774 +
775 + # Large batch: T=1000, E_per_gpu=1
776 + # Each GPU: T*E = 1000*1 = 1000 activations
777 + # Experts per GPU: 8
778 + # So num_activated_experts should be min(1000, 8) = 8 (all experts activated)
779 + large_ctx = ExecutionContext.from_single_request(
780 + num_tokens=1000, context_len=512, is_prefill=True
781 + )
782 + large_read = metrics.get_read_bytes_breakdown(large_ctx, per_gpu=True)
783 +
784 + # Weight reads should be similar (both activate all 8 experts per GPU)
785 + # But activation reads should differ (proportional to T*E)
786 + if "routed_up_gate_weights" in small_read:
787 + small_weight = small_read["routed_up_gate_weights"]
788 + large_weight = large_read["routed_up_gate_weights"]
789 +
790 + # Weight reads should be the same (both read all 8 experts)
791 + assert small_weight == large_weight
792 +
793 + # But input activation reads should scale with T*E
794 + small_input = small_read["routed_up_gate_input"]
795 + large_input = large_read["routed_up_gate_input"]
796 + assert large_input == 100 * small_input # 1000/10 = 100x
797 +
798 +
799 +def test_unembed_per_gpu_with_tensor_parallelism():
800 + """Test unembed metrics with tensor parallelism - per_gpu vs global."""
801 + hf_config = Qwen3Config(
802 + hidden_size=4096,
803 + vocab_size=128000,
804 + )
805 +
806 + # Test with TP=8
807 + vllm_config = create_mock_vllm_config(hf_config, tensor_parallel_size=8)
808 + metrics = UnembedMetrics.from_vllm_config(vllm_config)
809 +
810 + ctx = ExecutionContext.from_single_request(
811 + num_tokens=100, context_len=512, is_prefill=True
812 + )
813 +
814 + # Get global and per-gpu metrics
815 + global_flops = metrics.get_num_flops(ctx, per_gpu=False)
816 + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True)
817 +
818 + # With TP=8, vocab is divided by 8, so global should be 8x per-gpu
819 + assert global_flops == 8 * per_gpu_flops
820 +
821 + # For read bytes, weight reads scale with TP but input reads don't (replicated)
822 + global_read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=False)
823 + per_gpu_read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=True)
824 +
825 + # Input reads should be the same (replicated across TP ranks)
826 + assert global_read_breakdown["input"] == per_gpu_read_breakdown["input"]
827 +
828 + # Weight reads should scale 8x (divided by TP)
829 + assert global_read_breakdown["weight"] == 8 * per_gpu_read_breakdown["weight"]
830 +
831 +
832 +def test_model_metrics_per_gpu_aggregation():
833 + """Test ModelMetrics correctly aggregates per_gpu metrics across components."""
834 + hf_config = Qwen3Config(
835 + hidden_size=2048,
836 + num_attention_heads=16,
837 + num_hidden_layers=12,
838 + vocab_size=32000,
839 + intermediate_size=8192,
840 + )
841 +
842 + # Test with mixed parallelism: TP=2, PP=2
843 + vllm_config = create_mock_vllm_config(
844 + hf_config,
845 + tensor_parallel_size=2,
846 + pipeline_parallel_size=2,
847 + )
848 +
849 + model_metrics = ModelMetrics(vllm_config)
850 + ctx = ExecutionContext.from_single_request(
851 + num_tokens=100, context_len=512, is_prefill=True
852 + )
853 +
854 + # Get breakdowns for both modes
855 + per_gpu_breakdown = model_metrics.get_num_flops_breakdown(ctx, per_gpu=True)
856 + global_breakdown = model_metrics.get_num_flops_breakdown(ctx, per_gpu=False)
857 +
858 + # Verify breakdown sums match totals
859 + per_gpu_total = model_metrics.get_num_flops(ctx, per_gpu=True)
860 + global_total = model_metrics.get_num_flops(ctx, per_gpu=False)
861 +
862 + assert per_gpu_total == sum(per_gpu_breakdown.values())
863 + assert global_total == sum(global_breakdown.values())
864 +
865 + # Global should be larger than per-gpu due to parallelism
866 + assert global_total > per_gpu_total
867 +
868 + # With TP=2 and PP=2, the ratio depends on which parallelism applies to
869 + # which component but we can verify that global is reasonably larger
870 + ratio = global_total / per_gpu_total
871 + assert ratio > 1 # Should be between PP and TP*PP depending on component mix
872 +
873 +
874 +def test_attention_per_gpu_heads_not_evenly_divisible():
875 + """Test attention with heads not evenly divisible by TP."""
876 + hf_config = Qwen3Config(
877 + hidden_size=2048,
878 + num_attention_heads=17, # Not divisible by 4
879 + num_key_value_heads=5, # Not divisible by 4
880 + num_hidden_layers=8,
881 + )
882 +
883 + vllm_config = create_mock_vllm_config(hf_config, tensor_parallel_size=4)
884 + metrics = AttentionMetrics.from_vllm_config(vllm_config)
885 +
886 + ctx = ExecutionContext.from_single_request(
887 + num_tokens=64, context_len=256, is_prefill=True
888 + )
889 +
890 + # Should not crash and should handle max(1, ...) correctly
891 + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True)
892 + global_flops = metrics.get_num_flops(ctx, per_gpu=False)
893 +
894 + # Both should be positive
895 + assert per_gpu_flops > 0
896 + assert global_flops > 0
897 + assert global_flops > per_gpu_flops
vllm/config/observability.py
+3
@@ -64,6 +64,9 @@ class ObservabilityConfig:
64 module in the model and attach informations such as input/output shapes to
65 nvtx range markers. Noted that this doesn't work with CUDA graphs enabled."""
66
67 + enable_mfu_metrics: bool = False
68 + """Enable Model FLOPs Utilization (MFU) metrics."""
69 +
70 @cached_property
71 def collect_model_forward_time(self) -> bool:
72 """Whether to collect model forward time for the request."""
vllm/engine/arg_utils.py
+6
@@ -523,6 +523,7 @@ class EngineArgs:
523 enable_layerwise_nvtx_tracing: bool = (
524 ObservabilityConfig.enable_layerwise_nvtx_tracing
525 )
526 + enable_mfu_metrics: bool = ObservabilityConfig.enable_mfu_metrics
527 scheduling_policy: SchedulerPolicy = SchedulerConfig.policy
528 scheduler_cls: str | type[object] | None = SchedulerConfig.scheduler_cls
529
@@ -1042,6 +1043,10 @@ class EngineArgs:
1043 "--enable-layerwise-nvtx-tracing",
1044 **observability_kwargs["enable_layerwise_nvtx_tracing"],
1045 )
1046 + observability_group.add_argument(
1047 + "--enable-mfu-metrics",
1048 + **observability_kwargs["enable_mfu_metrics"],
1049 + )
1050
1051 # Scheduler arguments
1052 scheduler_kwargs = get_kwargs(SchedulerConfig)
@@ -1689,6 +1694,7 @@ class EngineArgs:
1694 kv_cache_metrics_sample=self.kv_cache_metrics_sample,
1695 cudagraph_metrics=self.cudagraph_metrics,
1696 enable_layerwise_nvtx_tracing=self.enable_layerwise_nvtx_tracing,
1697 + enable_mfu_metrics=self.enable_mfu_metrics,
1698 )
1699
1700 # Compilation config overrides
vllm/envs.py
+5
@@ -244,6 +244,7 @@ if TYPE_CHECKING:
244 VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
245 VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
246 VLLM_USE_V2_MODEL_RUNNER: bool = False
247 + VLLM_DEBUG_MFU_METRICS: bool = False
248
249
250 def get_default_cache_root():
@@ -1565,6 +1566,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1566 "VLLM_USE_V2_MODEL_RUNNER": lambda: bool(
1567 int(os.getenv("VLLM_USE_V2_MODEL_RUNNER", "0"))
1568 ),
1569 + # Debug logging for --enable-mfu-metrics
1570 + "VLLM_DEBUG_MFU_METRICS": lambda: bool(
1571 + int(os.getenv("VLLM_DEBUG_MFU_METRICS", "0"))
1572 + ),
1573 }
1574
1575 # --8<-- [end:env-vars-definition]
vllm/v1/core/sched/scheduler.py
+12 -1
@@ -43,6 +43,7 @@ from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_qu
43 from vllm.v1.core.sched.utils import check_stop, remove_all
44 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
45 from vllm.v1.kv_cache_interface import KVCacheConfig
46 +from vllm.v1.metrics.perf import ModelMetrics, PerfStats
47 from vllm.v1.metrics.stats import (
48 PrefixCacheStats,
49 SchedulerStats,
@@ -219,6 +220,10 @@ class Scheduler(SchedulerInterface):
220 self.use_pp = self.parallel_config.pipeline_parallel_size > 1
221 self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
222
223 + self.perf_metrics: ModelMetrics | None = None
224 + if self.log_stats and vllm_config.observability_config.enable_mfu_metrics:
225 + self.perf_metrics = ModelMetrics(vllm_config)
226 +
227 def schedule(self) -> SchedulerOutput:
228 # NOTE(woosuk) on the scheduling algorithm:
229 # There's no "decoding phase" nor "prefill phase" in the scheduler.
@@ -1066,6 +1071,10 @@ class Scheduler(SchedulerInterface):
1071 kv_connector_output = model_runner_output.kv_connector_output
1072 cudagraph_stats = model_runner_output.cudagraph_stats
1073
1074 + perf_stats: PerfStats | None = None
1075 + if self.perf_metrics and self.perf_metrics.is_enabled():
1076 + perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output)
1077 +
1078 outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
1079 spec_decoding_stats: SpecDecodingStats | None = None
1080 kv_connector_stats: KVConnectorStats | None = (
@@ -1262,7 +1271,7 @@ class Scheduler(SchedulerInterface):
1271
1272 if (
1273 stats := self.make_stats(
1265 - spec_decoding_stats, kv_connector_stats, cudagraph_stats
1274 + spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats
1275 )
1276 ) is not None:
1277 # Return stats to only one of the front-ends.
@@ -1485,6 +1494,7 @@ class Scheduler(SchedulerInterface):
1494 spec_decoding_stats: SpecDecodingStats | None = None,
1495 kv_connector_stats: KVConnectorStats | None = None,
1496 cudagraph_stats: CUDAGraphStat | None = None,
1497 + perf_stats: PerfStats | None = None,
1498 ) -> SchedulerStats | None:
1499 if not self.log_stats:
1500 return None
@@ -1510,6 +1520,7 @@ class Scheduler(SchedulerInterface):
1520 spec_decoding_stats=spec_stats,
1521 kv_connector_stats=connector_stats_payload,
1522 cudagraph_stats=cudagraph_stats,
1523 + perf_stats=perf_stats,
1524 )
1525
1526 def make_spec_decoding_stats(
vllm/v1/metrics/loggers.py
+16 -1
@@ -19,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
19 from vllm.logger import init_logger
20 from vllm.plugins import STAT_LOGGER_PLUGINS_GROUP, load_plugins_by_group
21 from vllm.v1.engine import FinishReason
22 +from vllm.v1.metrics.perf import PerfMetricsLogging
23 from vllm.v1.metrics.prometheus import unregister_vllm_metrics
24 from vllm.v1.metrics.stats import (
25 CachingMetrics,
@@ -118,6 +119,9 @@ class LoggingStatLogger(StatLoggerBase):
119 self.engine_is_idle = False
120 self.aggregated = False
121
122 + if self._enable_perf_stats():
123 + self.perf_metrics_logging = PerfMetricsLogging(vllm_config)
124 +
125 def _reset(self, now):
126 self.last_log_time = now
127
@@ -127,6 +131,9 @@ class LoggingStatLogger(StatLoggerBase):
131 self.num_corrupted_reqs: int = 0
132 self.num_preemptions: int = 0
133
134 + def _enable_perf_stats(self) -> bool:
135 + return self.vllm_config.observability_config.enable_mfu_metrics
136 +
137 def _track_iteration_stats(self, iteration_stats: IterationStats):
138 # Save tracked stats for token counters.
139 self.num_prompt_tokens += iteration_stats.num_prompt_tokens
@@ -175,6 +182,8 @@ class LoggingStatLogger(StatLoggerBase):
182 self.cudagraph_logging.observe(scheduler_stats.cudagraph_stats)
183 if not self.aggregated:
184 self.last_scheduler_stats = scheduler_stats
185 + if (perf_stats := scheduler_stats.perf_stats) and self._enable_perf_stats():
186 + self.perf_metrics_logging.observe(perf_stats)
187 if mm_cache_stats:
188 self.mm_caching_metrics.observe(mm_cache_stats)
189
@@ -211,7 +220,7 @@ class LoggingStatLogger(StatLoggerBase):
220 "Running: %d reqs",
221 "Waiting: %d reqs",
222 ]
214 - log_args = [
223 + log_args: list[int | float | str] = [
224 self.last_prompt_throughput,
225 self.last_generation_throughput,
226 self.last_scheduler_stats.num_running_reqs,
@@ -254,6 +263,8 @@ class LoggingStatLogger(StatLoggerBase):
263 self.kv_connector_logging.log(log_fn=log_fn)
264 if self.cudagraph_logging is not None:
265 self.cudagraph_logging.log(log_fn=log_fn)
266 + if self._enable_perf_stats():
267 + self.perf_metrics_logging.log(log_fn=log_fn, log_prefix=self.log_prefix)
268
269 def log_engine_initialized(self):
270 if self.vllm_config.cache_config.num_gpu_blocks:
@@ -282,6 +293,10 @@ class AggregatedLoggingStatLogger(LoggingStatLogger, AggregateStatLoggerBase):
293 def log_prefix(self):
294 return "{} Engines Aggregated: ".format(len(self.engine_indexes))
295
296 + def _enable_perf_stats(self) -> bool:
297 + # Adding per_gpu perf stats across engines can lead to misleading numbers.
298 + return False
299 +
300 def record(
301 self,
302 scheduler_stats: SchedulerStats | None,
vllm/v1/metrics/perf.py new
+1244
@@ -0,0 +1,1244 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +Analytic flops/memory estimation module for transformer components,
6 +to help derive MFU (Model Flops Utilization) stats for a running model.
7 +"""
8 +
9 +import json
10 +import time
11 +from abc import ABC, abstractmethod
12 +from collections.abc import Iterable
13 +from dataclasses import asdict, dataclass
14 +from typing import Any, Protocol
15 +
16 +import torch
17 +from pydantic import BaseModel, Field, ValidationError, model_validator
18 +from typing_extensions import Self
19 +
20 +import vllm.envs as envs
21 +from vllm.config import VllmConfig
22 +from vllm.logger import init_logger
23 +from vllm.utils.torch_utils import (
24 + STR_DTYPE_TO_TORCH_DTYPE,
25 + get_dtype_size,
26 + get_kv_cache_torch_dtype,
27 +)
28 +from vllm.v1.core.sched.output import SchedulerOutput
29 +
30 +logger = init_logger(__name__)
31 +
32 +
33 +class InvalidComponent(Exception):
34 + """
35 + Custom exception to indicate that a certain ComponentMetric is not
36 + applicable to the given VllmConfig.
37 + """
38 +
39 + pass
40 +
41 +
42 +#### Basic Data Types ####
43 +
44 +
45 +@dataclass
46 +class DebugPerfStats:
47 + ## Stats for debugging the metrics calculation
48 + calc_duration: float = 0.0 # time spent calculating these stats
49 + num_prefill_requests: int = 0
50 + num_decode_requests: int = 0
51 + context_breakdown: dict[str, int] | None = None
52 + num_flops_per_gpu_breakdown: dict[str, int] | None = None
53 + num_read_bytes_per_gpu_breakdown: dict[str, int] | None = None
54 + num_write_bytes_per_gpu_breakdown: dict[str, int] | None = None
55 +
56 +
57 +@dataclass
58 +class PerfStats:
59 + num_flops_per_gpu: int = 0
60 + num_read_bytes_per_gpu: int = 0
61 + num_write_bytes_per_gpu: int = 0
62 + debug_stats: DebugPerfStats | None = None
63 +
64 +
65 +@dataclass
66 +class ExecutionContext:
67 + """
68 + Represents an execution context for a batch of requests.
69 +
70 + This class aggregates statistics across multiple requests in a batch,
71 + separately tracking prefill and decode phases.
72 +
73 + Example)
74 + - Batch with one full prefill (2048 tokens) and one decode (1 token, 8192 context):
75 + ctx = ExecutionContext()
76 + ctx.add(2048, 2048, is_prefill=True)
77 + ctx.add(1, 8192, is_prefill=False)
78 + """
79 +
80 + # Prefill phase statistics
81 + num_prefill_requests: int = 0
82 + prefill_num_tokens: int = 0 # sum of num_tokens for prefill requests
83 + prefill_context_len: int = 0 # sum of context_len for prefill requests
84 + prefill_token_context_product: int = 0 # sum of (num_tokens * context_len)
85 +
86 + # Decode phase statistics
87 + num_decode_requests: int = 0
88 + decode_num_tokens: int = 0 # sum of num_tokens for decode requests
89 + decode_context_len: int = 0 # sum of context_len for decode requests
90 + decode_token_context_product: int = 0 # sum of (num_tokens * context_len)
91 +
92 + def add(self, num_tokens: int, context_len: int, is_prefill: bool) -> None:
93 + """Add a single request's statistics to this batch context."""
94 + if is_prefill:
95 + self.num_prefill_requests += 1
96 + self.prefill_num_tokens += num_tokens
97 + self.prefill_context_len += context_len
98 + self.prefill_token_context_product += num_tokens * context_len
99 + else:
100 + self.num_decode_requests += 1
101 + self.decode_num_tokens += num_tokens
102 + self.decode_context_len += context_len
103 + self.decode_token_context_product += num_tokens * context_len
104 +
105 + def total_num_tokens(self) -> int:
106 + """Total number of tokens across all requests in the batch."""
107 + return self.prefill_num_tokens + self.decode_num_tokens
108 +
109 + def total_token_context_product(self) -> int:
110 + """Total sum of (num_tokens * context_len) across all requests."""
111 + return self.prefill_token_context_product + self.decode_token_context_product
112 +
113 + @classmethod
114 + def from_single_request(
115 + cls, num_tokens: int, context_len: int, is_prefill: bool
116 + ) -> "ExecutionContext":
117 + """Create an ExecutionContext from a single request.
118 +
119 + This is a convenience method primarily for testing.
120 + """
121 + ctx = cls()
122 + ctx.add(num_tokens, context_len, is_prefill)
123 + return ctx
124 +
125 +
126 +class ParsedArgs:
127 + """
128 + Syntactic sugar so that Parsers can use dot notations
129 + to access/update the parsed arguments.
130 +
131 + e.g.)
132 + args = ParsedArgs()
133 + args.x = 3
134 + args.y = args.x + 1
135 + """
136 +
137 + def __getattr__(self, name: str) -> Any:
138 + raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")
139 +
140 + def __setattr__(self, name: str, value: Any) -> None:
141 + object.__setattr__(self, name, value)
142 +
143 + def model_dump(self) -> dict[str, Any]:
144 + return vars(self).copy()
145 +
146 +
147 +#### Abstract ####
148 +
149 +
150 +class Parser(Protocol):
151 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
152 + """
153 + Parse the vllm config and update the current ParsedArgs and pass it on.
154 + If the parser isn't applicable to the vllm_config, it will do nothing.
155 + """
156 + ...
157 +
158 +
159 +class ParserChain:
160 + """
161 + Applies chain of parser in a sequential order.
162 + Later parsers might overwrite results from previous parsers,
163 + so parsers should be chained in the appropriate order if they
164 + are not mutually exclusive.
165 + """
166 +
167 + def __init__(self, *parsers: Parser) -> None:
168 + self.parsers = list(parsers)
169 +
170 + def add_parser(self, parser: Parser) -> None:
171 + self.parsers.append(parser)
172 +
173 + def parse(self, vllm_config: VllmConfig) -> ParsedArgs:
174 + args = ParsedArgs()
175 + for parser in self.parsers:
176 + args = parser.parse(args, vllm_config)
177 + return args
178 +
179 +
180 +_COMPONENT_METRICS_REGISTRY: dict[str, type["ComponentMetrics"]] = {}
181 +
182 +
183 +class ComponentMetrics(BaseModel, ABC):
184 + """
185 + Each concrete ComponentMetrics class is associated with:
186 + - fields that are required for metric derivation
187 + (fields are specified/validated through pydantic model)
188 + - parser to parse VllmConfig into fields
189 + - metric methods that derive flops/bytes for a given execution context
190 + """
191 +
192 + @classmethod
193 + @abstractmethod
194 + def component_type(cls) -> str: ...
195 +
196 + @classmethod
197 + @abstractmethod
198 + def get_parser(cls) -> ParserChain:
199 + """
200 + Return a ParserChain that provides values for all required fields.
201 + The returned parser chain must populate ParsedArgs with values for every
202 + field defined on this ComponentMetrics class. Missing fields will cause
203 + a ValidationError when from_vllm_config() is called.
204 + See individual Parser docstrings for which args they provide, and field
205 + comments on ComponentMetrics subclasses for which parser provides each field.
206 + """
207 + ...
208 +
209 + def __init_subclass__(cls):
210 + _COMPONENT_METRICS_REGISTRY[cls.component_type()] = cls
211 +
212 + @classmethod
213 + def from_vllm_config(cls, vllm_config: VllmConfig) -> Self:
214 + """
215 + Instantiate this class from VllmConfig.
216 + Raises ValidationError if parsing fails.
217 + """
218 +
219 + parser = cls.get_parser()
220 + parsed_args = parser.parse(vllm_config)
221 + try:
222 + return cls.model_validate(parsed_args.model_dump())
223 + except ValidationError as e:
224 + raise InvalidComponent(f"Invalid {cls.component_type()} config: {e}") from e
225 +
226 + @classmethod
227 + def registered_metrics(cls) -> Iterable[type["ComponentMetrics"]]:
228 + return iter(_COMPONENT_METRICS_REGISTRY.values())
229 +
230 + @abstractmethod
231 + def get_num_flops_breakdown(
232 + self, ctx: ExecutionContext, per_gpu: bool = True
233 + ) -> dict[str, int]: ...
234 +
235 + @abstractmethod
236 + def get_read_bytes_breakdown(
237 + self, ctx: ExecutionContext, per_gpu: bool = True
238 + ) -> dict[str, int]: ...
239 +
240 + @abstractmethod
241 + def get_write_bytes_breakdown(
242 + self, ctx: ExecutionContext, per_gpu: bool = True
243 + ) -> dict[str, int]: ...
244 +
245 + def get_num_flops(self, ctx: ExecutionContext, per_gpu: bool = True) -> int:
246 + return sum(self.get_num_flops_breakdown(ctx, per_gpu).values())
247 +
248 + def get_read_bytes(self, ctx: ExecutionContext, per_gpu: bool = True) -> int:
249 + return sum(self.get_read_bytes_breakdown(ctx, per_gpu).values())
250 +
251 + def get_write_bytes(self, ctx: ExecutionContext, per_gpu: bool = True) -> int:
252 + return sum(self.get_write_bytes_breakdown(ctx, per_gpu).values())
253 +
254 +
255 +#### parsers ####
256 +
257 +
258 +class BaseConfigParser(Parser):
259 + """
260 + Parses base model configuration.
261 + Provides: vocab_size, hidden_size, num_attention_heads, num_hidden_layers,
262 + weight_byte_size, activation_byte_size, dp_size, tp_size, pp_size, enable_ep
263 + """
264 +
265 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
266 + model_config = vllm_config.model_config
267 +
268 + args.vocab_size = model_config.get_vocab_size()
269 + args.hidden_size = model_config.get_hidden_size()
270 + # NOTE: model_config.get_attention_heads() divide by TP
271 + # so we access field manually here to get total num_heads
272 + args.num_attention_heads = get_required(
273 + model_config.hf_text_config, "num_attention_heads"
274 + )
275 + args.num_hidden_layers = get_required(
276 + model_config.hf_text_config, "num_hidden_layers"
277 + )
278 +
279 + model_dtype = vllm_config.model_config.dtype
280 +
281 + if isinstance(model_dtype, torch.dtype):
282 + torch_dtype = model_dtype
283 + elif isinstance(model_dtype, str) and model_dtype in STR_DTYPE_TO_TORCH_DTYPE:
284 + torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
285 + else:
286 + # FIXME: handle this better
287 + logger.warning(
288 + "Unknown model_dtype %s, defaulting to bfloat16",
289 + model_dtype,
290 + )
291 + torch_dtype = torch.bfloat16
292 +
293 + args.weight_byte_size = get_dtype_size(torch_dtype)
294 +
295 + # FIXME: handle this better by parsing whether activations use
296 + # bf16, fp32, etc...
297 + args.activation_byte_size = 2
298 +
299 + args.dp_size = vllm_config.parallel_config.data_parallel_size
300 + args.tp_size = vllm_config.parallel_config.tensor_parallel_size
301 + args.pp_size = vllm_config.parallel_config.pipeline_parallel_size
302 + args.enable_ep = vllm_config.parallel_config.enable_expert_parallel
303 +
304 + return args
305 +
306 +
307 +#### Attention ####
308 +
309 +
310 +class BaseAttentionConfigParser(Parser):
311 + """
312 + Parses attention-specific configuration.
313 + Provides: num_key_value_heads, head_dim, cache_byte_size
314 + """
315 +
316 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
317 + model_config = vllm_config.model_config
318 +
319 + args.num_key_value_heads = model_config.get_total_num_kv_heads()
320 + args.head_dim = model_config.get_head_size()
321 +
322 + model_dtype = vllm_config.model_config.dtype
323 + cache_dtype = vllm_config.cache_config.cache_dtype
324 +
325 + kv_cache_torch_dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
326 + args.cache_byte_size = get_dtype_size(kv_cache_torch_dtype)
327 +
328 + return args
329 +
330 +
331 +class AttentionQuantizationConfigParser(Parser):
332 + """
333 + Parses quantization configuration for attention layers.
334 + Overrides: weight_byte_size
335 + """
336 +
337 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
338 + cfg = vllm_config.quant_config
339 +
340 + if cfg is None:
341 + return args
342 +
343 + quant_method = cfg.get_name()
344 + if quant_method in ["fp8", "fbgemm_fp8"]:
345 + # FIXME: This is a hacky coarse-grained fp8 quantization detection.
346 + # FIXME: These configs also have concept of "ignored layers" and we
347 + # need to solve the same problem as above.
348 + args.weight_byte_size = 1
349 + elif quant_method == "mxfp4":
350 + # FIXME: Also has "ignored layers" issue above
351 + args.weight_byte_size = 0.5
352 + else:
353 + # FIXME: Add more parsing logic for different quant methods.
354 + raise InvalidComponent
355 +
356 + return args
357 +
358 +
359 +class AttentionMetrics(ComponentMetrics):
360 + # From BaseConfigParser
361 + num_hidden_layers: int = Field(..., gt=0)
362 + hidden_size: int = Field(..., gt=0)
363 + num_attention_heads: int = Field(..., gt=0)
364 + activation_byte_size: int = Field(..., gt=0)
365 + tp_size: int = Field(..., gt=0)
366 + pp_size: int = Field(..., gt=0)
367 +
368 + # From BaseAttentionConfigParser
369 + num_key_value_heads: int = Field(..., gt=0)
370 + head_dim: int = Field(..., gt=0)
371 + cache_byte_size: int = Field(..., gt=0)
372 +
373 + # From BaseConfig Parser, overridden by AttentionQuantizationConfigParser
374 + weight_byte_size: int | float = Field(..., gt=0)
375 +
376 + # TODO: discern cases where we have mixture of different attention layer types
377 + # such as SWA, MLA, etc.
378 +
379 + @classmethod
380 + def component_type(cls) -> str:
381 + return "attn"
382 +
383 + @classmethod
384 + def get_parser(cls) -> ParserChain:
385 + return ParserChain(
386 + BaseConfigParser(),
387 + BaseAttentionConfigParser(),
388 + AttentionQuantizationConfigParser(),
389 + )
390 +
391 + def get_num_flops_breakdown(
392 + self, ctx: ExecutionContext, per_gpu: bool = True
393 + ) -> dict[str, int]:
394 + L, D, q, kv, d = (
395 + self.num_hidden_layers,
396 + self.hidden_size,
397 + self.num_attention_heads,
398 + self.num_key_value_heads,
399 + self.head_dim,
400 + )
401 + T = ctx.total_num_tokens()
402 + TC = ctx.total_token_context_product()
403 +
404 + if per_gpu:
405 + L //= self.pp_size
406 + # tensor parallel along heads
407 + q = max(1, q // self.tp_size)
408 + kv = max(1, kv // self.tp_size)
409 +
410 + return {
411 + "qkv_proj": 2 * T * D * (q + 2 * kv) * d * L,
412 + "attn_qk": 2 * q * TC * d * L,
413 + "attn_av": 2 * q * TC * d * L,
414 + "out_proj": 2 * T * D * q * d * L,
415 + }
416 +
417 + def get_read_bytes_breakdown(
418 + self, ctx: ExecutionContext, per_gpu: bool = True
419 + ) -> dict[str, int]:
420 + L, D, q, kv, d = (
421 + self.num_hidden_layers,
422 + self.hidden_size,
423 + self.num_attention_heads,
424 + self.num_key_value_heads,
425 + self.head_dim,
426 + )
427 + T = ctx.total_num_tokens()
428 +
429 + if per_gpu:
430 + L //= self.pp_size
431 + # tensor parallel along heads
432 + q = max(1, q // self.tp_size)
433 + kv = max(1, kv // self.tp_size)
434 +
435 + read_bytes = {}
436 +
437 + read_bytes["qkv_input"] = T * D * self.activation_byte_size * L
438 + read_bytes["qkv_weight"] = int(D * (q + 2 * kv) * d * self.weight_byte_size * L)
439 +
440 + # Attention input reads differ between prefill and decode
441 + # Prefill: read Q, K, V activations (all in activation_byte_size)
442 + if ctx.prefill_num_tokens > 0:
443 + read_bytes["attn_input"] = (
444 + (ctx.prefill_num_tokens * q + 2 * ctx.prefill_context_len * kv)
445 + * d
446 + * self.activation_byte_size
447 + * L
448 + )
449 +
450 + # Decode: read Q activations + read K, V from cache (in cache_byte_size)
451 + if ctx.decode_num_tokens > 0:
452 + read_bytes["attn_input"] = read_bytes.get("attn_input", 0) + (
453 + ctx.decode_num_tokens * q * d * self.activation_byte_size * L
454 + + 2 * ctx.decode_context_len * kv * d * self.cache_byte_size * L
455 + )
456 +
457 + read_bytes["out_input"] = T * q * d * self.activation_byte_size * L
458 + read_bytes["out_weight"] = int(q * d * D * self.weight_byte_size * L)
459 +
460 + return read_bytes
461 +
462 + def get_write_bytes_breakdown(
463 + self, ctx: ExecutionContext, per_gpu: bool = True
464 + ) -> dict[str, int]:
465 + """Calculate write memory traffic for attention layers."""
466 + L, D, q, kv, d = (
467 + self.num_hidden_layers,
468 + self.hidden_size,
469 + self.num_attention_heads,
470 + self.num_key_value_heads,
471 + self.head_dim,
472 + )
473 + T = ctx.total_num_tokens()
474 +
475 + if per_gpu:
476 + L //= self.pp_size
477 + # tensor parallel along heads
478 + q = max(1, q // self.tp_size)
479 + kv = max(1, kv // self.tp_size)
480 +
481 + return {
482 + "qkv_output": T * (q + 2 * kv) * d * self.activation_byte_size * L,
483 + "kv_cache": 2 * T * kv * d * self.cache_byte_size * L,
484 + "out_output": T * D * self.activation_byte_size * L,
485 + }
486 +
487 +
488 +#### Ffn ####
489 +
490 +
491 +class BaseFfnConfigParser(Parser):
492 + """
493 + Parses FFN and MoE configuration.
494 + Provides: intermediate_size, num_experts, num_experts_per_tok,
495 + moe_intermediate_size, num_shared_experts, num_moe_layers
496 + """
497 +
498 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
499 + cfg = vllm_config.model_config.hf_config
500 + if hasattr(cfg, "text_config") and cfg.text_config is not None:
501 + cfg = cfg.text_config
502 +
503 + args.intermediate_size = getattr(cfg, "intermediate_size", args.hidden_size * 4)
504 +
505 + # Try different naming conventions.
506 + args.num_experts = vllm_config.model_config.get_num_experts()
507 + args.num_experts_per_tok = getattr_from_list(
508 + cfg, ["num_experts_per_tok", "moe_topk"], 0
509 + )
510 + args.moe_intermediate_size = getattr_from_list(
511 + cfg, ["moe_intermediate_size", "intermediate_size"], 0
512 + )
513 + args.num_shared_experts = getattr_from_list(
514 + cfg, ["n_shared_experts", "num_shared_experts"], 0
515 + )
516 +
517 + is_moe = args.num_experts != 0
518 + # Assume all MoE layers by default
519 + args.num_moe_layers = args.num_hidden_layers if is_moe else 0
520 +
521 + return args
522 +
523 +
524 +class FfnParallelParser(Parser):
525 + """
526 + Parses FFN parallelism configuration.
527 +
528 + Provides: ffn_tp_size, ffn_ep_size
529 + """
530 +
531 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
532 + # NOTE: ffn tp_size does not equal the tp_size parameter directly.
533 + # e.g.) If we use DP2TP4, ffn will use TP8 (or EP8 if EP is enabled.)
534 + if args.enable_ep:
535 + ffn_tp_size, ffn_ep_size = 1, args.dp_size * args.tp_size
536 + else:
537 + ffn_tp_size, ffn_ep_size = args.dp_size * args.tp_size, 1
538 +
539 + args.ffn_tp_size = ffn_tp_size
540 + args.ffn_ep_size = ffn_ep_size
541 +
542 + return args
543 +
544 +
545 +class InterleaveMoeLayerStepParser(Parser):
546 + """
547 + Parses interleave_moe_layer_step field for models like Llama4.
548 +
549 + Overrides: num_moe_layers
550 + """
551 +
552 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
553 + cfg = vllm_config.model_config.hf_config
554 + if hasattr(cfg, "text_config") and cfg.text_config is not None:
555 + cfg = cfg.text_config
556 +
557 + if (
558 + hasattr(cfg, "interleave_moe_layer_step")
559 + and cfg.interleave_moe_layer_step > 0
560 + ):
561 + args.num_moe_layers = len(
562 + [
563 + layer
564 + for layer in range(args.num_hidden_layers)
565 + if (layer + 1) % cfg.interleave_moe_layer_step == 0
566 + ]
567 + )
568 +
569 + return args
570 +
571 +
572 +class MoeLayerFreqParser(Parser):
573 + """
574 + Parses moe_layer_freq and first_k_dense_replace fields for models like Deepseek.
575 +
576 + Overrides: num_moe_layers
577 + """
578 +
579 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
580 + cfg = vllm_config.model_config.hf_config
581 + if hasattr(cfg, "text_config") and cfg.text_config is not None:
582 + cfg = cfg.text_config
583 +
584 + if hasattr(cfg, "moe_layer_freq") and hasattr(cfg, "first_k_dense_replace"):
585 + args.num_moe_layers = len(
586 + [
587 + layer
588 + for layer in range(args.num_hidden_layers)
589 + if layer >= cfg.first_k_dense_replace
590 + and layer % cfg.moe_layer_freq == 0
591 + ]
592 + )
593 +
594 + return args
595 +
596 +
597 +class FfnQuantizationConfigParser(Parser):
598 + """
599 + Parses quantization configuration for FFN layers.
600 +
601 + Overrides: weight_byte_size
602 + """
603 +
604 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
605 + cfg = vllm_config.quant_config
606 +
607 + if cfg is None:
608 + return args
609 +
610 + quant_method = cfg.get_name()
611 + if quant_method in ["fp8", "fbgemm_fp8"]:
612 + # FIXME: This is a hacky coarse-grained fp8 quantization detection.
613 + # (there might be more quantization methods for fp8).
614 + # FIXME: These configs also have concept of "ignored layers" and we
615 + # need to solve the same problem as above.
616 + args.weight_byte_size = 1
617 + pass
618 + elif quant_method == "mxfp4":
619 + # FIXME: Also has "ignored layers" issue above
620 + args.weight_byte_size = 0.5
621 + else:
622 + # FIXME: Add more parsing logic for different quant methods.
623 + raise InvalidComponent
624 +
625 + return args
626 +
627 +
628 +class FfnMetrics(ComponentMetrics):
629 + # From BaseConfigParser
630 + num_hidden_layers: int = Field(..., gt=0)
631 + hidden_size: int = Field(..., gt=0)
632 + activation_byte_size: int = Field(..., gt=0)
633 + pp_size: int = Field(..., gt=0)
634 +
635 + # From FfnParallelParser
636 + ffn_tp_size: int = Field(..., gt=0)
637 + ffn_ep_size: int = Field(..., gt=0)
638 +
639 + # From BaseFfnConfigParser
640 + intermediate_size: int = Field(..., gt=0)
641 + num_experts: int = Field(0)
642 + num_experts_per_tok: int = Field(1)
643 + moe_intermediate_size: int = Field(0)
644 + num_shared_experts: int = Field(0)
645 +
646 + # From BaseConfigParser, can be overridden InterleaveMoeLayerStep or MoeLayerFreq
647 + num_moe_layers: int = Field(..., ge=0)
648 +
649 + # FIXME: might have to make this more granular
650 + # (i.e. dense_weight_byte_size, moe_routed_weight_byte_size,
651 + # moe_shared_weight_byte_size)
652 + # since it can differ from byte size of other components (e.g. attn)
653 + # and can differ even from each other.
654 +
655 + # From BaseConfigParser, can be overridden by FfnQuantizationConfigParser
656 + weight_byte_size: int | float = Field(..., gt=0)
657 +
658 + @model_validator(mode="after")
659 + def validate_moe_fields(self) -> Self:
660 + """Validate that MoE-related fields are properly set when num_moe_layers > 0."""
661 + if self.num_moe_layers > 0:
662 + assert self.num_experts, f"{self.num_experts=}"
663 + assert self.num_experts_per_tok, f"{self.num_experts_per_tok=}"
664 + assert self.moe_intermediate_size, f"{self.moe_intermediate_size=}"
665 + return self
666 +
667 + @classmethod
668 + def component_type(cls) -> str:
669 + return "ffn"
670 +
671 + @classmethod
672 + def get_parser(cls) -> ParserChain:
673 + return ParserChain(
674 + BaseConfigParser(),
675 + FfnParallelParser(),
676 + BaseFfnConfigParser(),
677 + InterleaveMoeLayerStepParser(),
678 + MoeLayerFreqParser(),
679 + FfnQuantizationConfigParser(),
680 + )
681 +
682 + def get_num_flops_breakdown(
683 + self, ctx: ExecutionContext, per_gpu: bool = True
684 + ) -> dict[str, int]:
685 + """Calculate flops breakdown for FFN layers."""
686 + L, D, DI = self.num_hidden_layers, self.hidden_size, self.intermediate_size
687 + Lm, E, MI, S = (
688 + self.num_moe_layers,
689 + self.num_experts_per_tok,
690 + self.moe_intermediate_size,
691 + self.num_shared_experts,
692 + )
693 + T = ctx.total_num_tokens()
694 +
695 + Ld = L - Lm
696 +
697 + num_activated_tokens = T * E if E else 0
698 +
699 + if per_gpu:
700 + Ld //= self.pp_size
701 + Lm //= self.pp_size
702 +
703 + DI //= self.ffn_tp_size
704 + if MI is not None:
705 + MI //= self.ffn_tp_size
706 + if E:
707 + num_activated_tokens //= self.ffn_ep_size
708 +
709 + flops = {}
710 +
711 + # Dense FFN layers (SwiGLU: 3 linear layers: up, gate, down)
712 + if Ld:
713 + flops["dense_ffn"] = 2 * D * 3 * DI * T * Ld
714 +
715 + # MoE routed experts (each token activates E experts)
716 + if Lm and E:
717 + flops["routed_ffn"] = 2 * D * 3 * MI * num_activated_tokens * Lm
718 +
719 + # MoE shared experts (all S shared experts run for every token)
720 + if Lm and S:
721 + flops["shared_ffn"] = 2 * D * 3 * MI * S * T * Lm
722 +
723 + return flops
724 +
725 + def get_read_bytes_breakdown(
726 + self, ctx: ExecutionContext, per_gpu: bool = True
727 + ) -> dict[str, int]:
728 + """Calculate read memory traffic for FFN layers."""
729 + L, D, DI = self.num_hidden_layers, self.hidden_size, self.intermediate_size
730 + Lm, E, MI, S = (
731 + self.num_moe_layers,
732 + self.num_experts_per_tok,
733 + self.moe_intermediate_size,
734 + self.num_shared_experts,
735 + )
736 + T = ctx.total_num_tokens()
737 + num_experts = self.num_experts
738 +
739 + Ld = L - Lm
740 +
741 + num_activated_tokens = T * E if E else 0
742 +
743 + if per_gpu:
744 + Ld //= self.pp_size
745 + Lm //= self.pp_size
746 +
747 + DI //= self.ffn_tp_size
748 + if MI is not None:
749 + MI //= self.ffn_tp_size
750 + if E:
751 + num_activated_tokens //= self.ffn_ep_size
752 + if num_experts is not None:
753 + num_experts //= self.ffn_ep_size
754 +
755 + read_bytes = {}
756 +
757 + # Dense FFN layers (3 GEMMs: up, gate, down projections + SiLU activation)
758 + if Ld:
759 + read_bytes["dense_up_gate_input"] = int(
760 + T * D * self.activation_byte_size * Ld
761 + )
762 + read_bytes["dense_up_gate_weights"] = int(
763 + 2 * D * DI * self.weight_byte_size * Ld
764 + )
765 + read_bytes["dense_silu_input"] = int(
766 + 2 * T * DI * self.activation_byte_size * Ld
767 + )
768 + read_bytes["dense_down_input"] = int(
769 + T * DI * self.activation_byte_size * Ld
770 + )
771 + read_bytes["dense_down_weights"] = int(D * DI * self.weight_byte_size * Ld)
772 +
773 + if Lm:
774 + # MoE routed expert reads
775 + if E:
776 + # FIXME: Assume perfect load balancing for now.
777 + num_activated_experts = min(num_activated_tokens, num_experts)
778 +
779 + read_bytes["routed_up_gate_input"] = int(
780 + num_activated_tokens * D * self.activation_byte_size * Lm
781 + )
782 + read_bytes["routed_up_gate_weights"] = int(
783 + 2 * D * MI * num_activated_experts * self.weight_byte_size * Lm
784 + )
785 + read_bytes["routed_silu_input"] = int(
786 + 2 * num_activated_tokens * MI * self.activation_byte_size * Lm
787 + )
788 + read_bytes["routed_down_input"] = int(
789 + num_activated_tokens * MI * self.activation_byte_size * Lm
790 + )
791 + read_bytes["routed_down_weights"] = int(
792 + D * MI * num_activated_experts * self.weight_byte_size * Lm
793 + )
794 +
795 + # MoE shared expert reads
796 + if S:
797 + read_bytes["shared_up_gate_input"] = int(
798 + T * D * self.activation_byte_size * Lm
799 + )
800 + read_bytes["shared_up_gate_weights"] = int(
801 + 2 * D * MI * S * self.weight_byte_size * Lm
802 + )
803 + read_bytes["shared_silu_input"] = int(
804 + 2 * T * MI * S * self.activation_byte_size * Lm
805 + )
806 + read_bytes["shared_down_input"] = int(
807 + T * MI * self.activation_byte_size * Lm
808 + )
809 + read_bytes["shared_down_weights"] = int(
810 + D * MI * S * self.weight_byte_size * Lm
811 + )
812 +
813 + return read_bytes
814 +
815 + def get_write_bytes_breakdown(
816 + self, ctx: ExecutionContext, per_gpu: bool = True
817 + ) -> dict[str, int]:
818 + """Calculate write memory traffic for FFN layers."""
819 + L, D, DI = self.num_hidden_layers, self.hidden_size, self.intermediate_size
820 + Lm, E, MI, S = (
821 + self.num_moe_layers,
822 + self.num_experts_per_tok,
823 + self.moe_intermediate_size,
824 + self.num_shared_experts,
825 + )
826 + T = ctx.total_num_tokens()
827 +
828 + Ld = L - Lm
829 +
830 + num_activated_tokens = T * E if E else 0
831 +
832 + if per_gpu:
833 + Ld //= self.pp_size
834 + Lm //= self.pp_size
835 +
836 + DI //= self.ffn_tp_size
837 + if MI is not None:
838 + MI //= self.ffn_tp_size
839 + if E:
840 + num_activated_tokens //= self.ffn_ep_size
841 +
842 + write_bytes = {}
843 +
844 + # Dense FFN layers
845 + if Ld:
846 + write_bytes["dense_up_gate_output"] = int(
847 + 2 * T * DI * self.activation_byte_size * Ld
848 + )
849 + write_bytes["dense_silu_output"] = int(
850 + T * DI * self.activation_byte_size * Ld
851 + )
852 + write_bytes["dense_down_output"] = int(
853 + T * D * self.activation_byte_size * Ld
854 + )
855 +
856 + # MoE outputs
857 + if Lm:
858 + if E:
859 + write_bytes["routed_up_gate_output"] = int(
860 + 2 * num_activated_tokens * MI * self.activation_byte_size * Lm
861 + )
862 + write_bytes["routed_silu_output"] = int(
863 + num_activated_tokens * MI * self.activation_byte_size * Lm
864 + )
865 + write_bytes["routed_down_output"] = int(
866 + num_activated_tokens * D * self.activation_byte_size * Lm
867 + )
868 + if S:
869 + write_bytes["shared_up_gate_output"] = int(
870 + 2 * T * S * MI * self.activation_byte_size * Lm
871 + )
872 + write_bytes["shared_silu_output"] = int(
873 + T * S * MI * self.activation_byte_size * Lm
874 + )
875 + write_bytes["shared_down_output"] = int(
876 + T * S * D * self.activation_byte_size * Lm
877 + )
878 +
879 + return write_bytes
880 +
881 +
882 +#### Unembed ####
883 +
884 +
885 +class UnembedMetrics(ComponentMetrics):
886 + # From BaseConfigParser
887 + hidden_size: int = Field(..., gt=0)
888 + vocab_size: int = Field(..., gt=0)
889 + weight_byte_size: int = Field(..., gt=0)
890 + activation_byte_size: int = Field(..., gt=0)
891 +
892 + tp_size: int
893 +
894 + @classmethod
895 + def component_type(cls) -> str:
896 + return "unembed"
897 +
898 + @classmethod
899 + def get_parser(cls) -> ParserChain:
900 + return ParserChain(
901 + BaseConfigParser(),
902 + )
903 +
904 + def get_num_flops_breakdown(
905 + self, ctx: ExecutionContext, per_gpu: bool = True
906 + ) -> dict[str, int]:
907 + """Calculate flops breakdown for unembedding layer."""
908 + D, V = self.hidden_size, self.vocab_size
909 + T = ctx.total_num_tokens()
910 +
911 + if per_gpu:
912 + V //= self.tp_size
913 +
914 + return {
915 + "unembed": 2 * T * D * V,
916 + }
917 +
918 + def get_read_bytes_breakdown(
919 + self, ctx: ExecutionContext, per_gpu: bool = True
920 + ) -> dict[str, int]:
921 + """Calculate read memory traffic for unembedding layer."""
922 + D, V = self.hidden_size, self.vocab_size
923 + T = ctx.total_num_tokens()
924 +
925 + if per_gpu:
926 + V //= self.tp_size
927 +
928 + return {
929 + "input": T * D * self.activation_byte_size,
930 + "weight": D * V * self.weight_byte_size,
931 + }
932 +
933 + def get_write_bytes_breakdown(
934 + self, ctx: ExecutionContext, per_gpu: bool = True
935 + ) -> dict[str, int]:
936 + """Calculate write memory traffic for unembedding layer."""
937 + V = self.vocab_size
938 + T = ctx.total_num_tokens()
939 +
940 + if per_gpu:
941 + V //= self.tp_size
942 +
943 + return {
944 + "output": T * V * self.activation_byte_size,
945 + }
946 +
947 +
948 +#### ModelMetrics ####
949 +
950 +
951 +class ModelMetrics:
952 + def __init__(self, vllm_config: VllmConfig) -> None:
953 + """
954 + Parse vllm_config to instantiate metrics for each component.
955 + is_enabled() will return False if no component metrics could be instantiated.
956 + """
957 +
958 + self.vllm_config = vllm_config
959 +
960 + self.metrics: list[ComponentMetrics] = []
961 + for metric_cls in ComponentMetrics.registered_metrics():
962 + try:
963 + metric = metric_cls.from_vllm_config(vllm_config)
964 + self.metrics.append(metric)
965 + logger.info(
966 + "Instantiated ComponentMetrics [%s] with (%s)",
967 + metric.component_type(),
968 + str(metric),
969 + )
970 + except InvalidComponent as e:
971 + logger.debug(
972 + "Failed to instantiate %s from %s",
973 + metric_cls.component_type(),
974 + str(e),
975 + )
976 +
977 + def is_enabled(self) -> bool:
978 + return len(self.metrics) > 0
979 +
980 + def get_num_flops(self, ctx: ExecutionContext, per_gpu: bool = True) -> int:
981 + return sum(metric.get_num_flops(ctx, per_gpu) for metric in self.metrics)
982 +
983 + def get_read_bytes(self, ctx: ExecutionContext, per_gpu: bool = True) -> int:
984 + return sum(metric.get_read_bytes(ctx, per_gpu) for metric in self.metrics)
985 +
986 + def get_write_bytes(self, ctx: ExecutionContext, per_gpu: bool = True) -> int:
987 + return sum(metric.get_write_bytes(ctx, per_gpu) for metric in self.metrics)
988 +
989 + def get_num_flops_breakdown(
990 + self, ctx: ExecutionContext, per_gpu: bool = True
991 + ) -> dict[str, int]:
992 + total = {}
993 + for metric in self.metrics:
994 + breakdown = metric.get_num_flops_breakdown(ctx, per_gpu)
995 + component = metric.component_type()
996 + prefixed = {f"{component}.{key}": val for key, val in breakdown.items()}
997 + total.update(prefixed)
998 + return total
999 +
1000 + def get_read_bytes_breakdown(
1001 + self, ctx: ExecutionContext, per_gpu: bool = True
1002 + ) -> dict[str, int]:
1003 + total = {}
1004 + for metric in self.metrics:
1005 + breakdown = metric.get_read_bytes_breakdown(ctx, per_gpu)
1006 + component = metric.component_type()
1007 + prefixed = {f"{component}.{key}": val for key, val in breakdown.items()}
1008 + total.update(prefixed)
1009 + return total
1010 +
1011 + def get_write_bytes_breakdown(
1012 + self, ctx: ExecutionContext, per_gpu: bool = True
1013 + ) -> dict[str, int]:
1014 + total = {}
1015 + for metric in self.metrics:
1016 + breakdown = metric.get_write_bytes_breakdown(ctx, per_gpu)
1017 + component = metric.component_type()
1018 + prefixed = {f"{component}.{key}": val for key, val in breakdown.items()}
1019 + total.update(prefixed)
1020 + return total
1021 +
1022 + def get_step_perf_stats_per_gpu(
1023 + self, scheduler_output: SchedulerOutput
1024 + ) -> PerfStats:
1025 + """
1026 + Calculate perf stats for the current step based on scheduled tokens.
1027 + """
1028 +
1029 + t0 = time.monotonic()
1030 +
1031 + # Build a single batch context
1032 + ctx = ExecutionContext()
1033 +
1034 + # Process new requests (these are in prefill phase)
1035 + for new_req in scheduler_output.scheduled_new_reqs:
1036 + req_id = new_req.req_id
1037 + num_tokens = scheduler_output.num_scheduled_tokens.get(req_id, 0)
1038 + if num_tokens == 0:
1039 + continue
1040 +
1041 + # For new requests, context_len = num_computed_tokens + num_tokens
1042 + # num_computed_tokens represents previously computed tokens in the sequence
1043 + context_len = new_req.num_computed_tokens + num_tokens
1044 + ctx.add(num_tokens, context_len, is_prefill=True)
1045 +
1046 + # Process cached requests (continuing requests)
1047 + cached_reqs = scheduler_output.scheduled_cached_reqs
1048 + for i, req_id in enumerate(cached_reqs.req_ids):
1049 + num_tokens = scheduler_output.num_scheduled_tokens.get(req_id, 0)
1050 + if num_tokens == 0:
1051 + continue
1052 +
1053 + # For cached requests, we have the current num_computed_tokens
1054 + num_computed_tokens = cached_reqs.num_computed_tokens[i]
1055 + context_len = num_computed_tokens + num_tokens
1056 +
1057 + # Cached requests are typically in decode phase (num_tokens == 1)
1058 + # unless they're doing chunked prefill (num_tokens > 1)
1059 + is_prefill = num_tokens > 1
1060 + ctx.add(num_tokens, context_len, is_prefill)
1061 +
1062 + num_flops_breakdown = self.get_num_flops_breakdown(ctx, True)
1063 + read_bytes_breakdown = self.get_read_bytes_breakdown(ctx, True)
1064 + write_bytes_breakdown = self.get_write_bytes_breakdown(ctx, True)
1065 + perf_stats = PerfStats(
1066 + sum(num_flops_breakdown.values()),
1067 + sum(read_bytes_breakdown.values()),
1068 + sum(write_bytes_breakdown.values()),
1069 + )
1070 +
1071 + if envs.VLLM_DEBUG_MFU_METRICS:
1072 + perf_stats.debug_stats = DebugPerfStats(
1073 + time.monotonic() - t0,
1074 + ctx.num_prefill_requests,
1075 + ctx.num_decode_requests,
1076 + asdict(ctx),
1077 + num_flops_breakdown,
1078 + read_bytes_breakdown,
1079 + write_bytes_breakdown,
1080 + )
1081 +
1082 + return perf_stats
1083 +
1084 +
1085 +#### Logging ####
1086 +
1087 +
1088 +class PerfMetricsDebugLogging:
1089 + def __init__(self):
1090 + self.reset()
1091 +
1092 + def reset(self):
1093 + self.total_calc_duration: float = 0.0
1094 + self.total_num_prefill_requests: int = 0
1095 + self.total_num_decode_requests: int = 0
1096 + self.total_num_batches: int = 0
1097 + self.total_context_breakdown: dict[str, int] = {}
1098 + self.total_num_flops_per_gpu_breakdown: dict[str, int] = {}
1099 + self.total_read_bytes_per_gpu_breakdown: dict[str, int] = {}
1100 + self.total_write_bytes_per_gpu_breakdown: dict[str, int] = {}
1101 +
1102 + def observe(self, debug_stats: DebugPerfStats) -> None:
1103 + self.total_calc_duration += debug_stats.calc_duration
1104 + self.total_num_prefill_requests += debug_stats.num_prefill_requests
1105 + self.total_num_decode_requests += debug_stats.num_decode_requests
1106 + self.total_num_batches += 1
1107 +
1108 + for dst, src in zip(
1109 + [
1110 + self.total_context_breakdown,
1111 + self.total_num_flops_per_gpu_breakdown,
1112 + self.total_read_bytes_per_gpu_breakdown,
1113 + self.total_write_bytes_per_gpu_breakdown,
1114 + ],
1115 + [
1116 + debug_stats.context_breakdown,
1117 + debug_stats.num_flops_per_gpu_breakdown,
1118 + debug_stats.num_read_bytes_per_gpu_breakdown,
1119 + debug_stats.num_write_bytes_per_gpu_breakdown,
1120 + ],
1121 + ):
1122 + assert isinstance(src, dict)
1123 + for key, val in src.items():
1124 + dst[key] = dst.get(key, 0) + val
1125 +
1126 + def log(self, log_fn, log_prefix: str, delta_time: float):
1127 + # pretty print breakdowns
1128 + total_num_flops_per_gpu_breakdown = {
1129 + k: f"{v / 1e12:.1f}TF"
1130 + for k, v in self.total_num_flops_per_gpu_breakdown.items()
1131 + }
1132 + total_read_bytes_per_gpu_breakdown = {
1133 + k: f"{v / 1e9:.1f}GB"
1134 + for k, v in self.total_read_bytes_per_gpu_breakdown.items()
1135 + }
1136 + total_write_bytes_per_gpu_breakdown = {
1137 + k: f"{v / 1e9:.1f}GB"
1138 + for k, v in self.total_write_bytes_per_gpu_breakdown.items()
1139 + }
1140 +
1141 + logger.debug(
1142 + "%sMFU details: %s",
1143 + log_prefix,
1144 + json.dumps(
1145 + {
1146 + "prefill_reqs": self.total_num_prefill_requests,
1147 + "decode_reqs": self.total_num_decode_requests,
1148 + "num_batches": self.total_num_batches,
1149 + "context_breakdown": self.total_context_breakdown,
1150 + "flops_breakdown": total_num_flops_per_gpu_breakdown,
1151 + "num_read_bytes_breakdown": total_read_bytes_per_gpu_breakdown,
1152 + "num_write_bytes_breakdown": (total_write_bytes_per_gpu_breakdown),
1153 + "duration": f"{delta_time:.1f}s",
1154 + "mfu_calc_overhead": (
1155 + f"{self.total_calc_duration / delta_time:.1%}"
1156 + ),
1157 + },
1158 + indent=2,
1159 + ),
1160 + )
1161 +
1162 +
1163 +class PerfMetricsLogging:
1164 + def __init__(self, vllm_config: VllmConfig):
1165 + self.vllm_config = vllm_config
1166 + self.pp_size = vllm_config.parallel_config.pipeline_parallel_size
1167 +
1168 + self.debug_logging: PerfMetricsDebugLogging | None = None
1169 + if envs.VLLM_DEBUG_MFU_METRICS:
1170 + self.debug_logging = PerfMetricsDebugLogging()
1171 +
1172 + self.reset()
1173 +
1174 + def reset(self):
1175 + self.last_log_time = time.monotonic()
1176 +
1177 + self.total_num_flops_per_gpu: int = 0
1178 + self.total_read_bytes_per_gpu: int = 0
1179 + self.total_write_bytes_per_gpu: int = 0
1180 +
1181 + if self.debug_logging:
1182 + self.debug_logging.reset()
1183 +
1184 + def observe(self, perf_stats: PerfStats) -> None:
1185 + self.total_num_flops_per_gpu += perf_stats.num_flops_per_gpu
1186 + self.total_read_bytes_per_gpu += perf_stats.num_read_bytes_per_gpu
1187 + self.total_write_bytes_per_gpu += perf_stats.num_write_bytes_per_gpu
1188 +
1189 + if self.debug_logging:
1190 + assert perf_stats.debug_stats is not None
1191 + self.debug_logging.observe(perf_stats.debug_stats)
1192 +
1193 + def log(self, log_fn=logger.info, log_prefix: str = "") -> None:
1194 + if not (
1195 + self.total_num_flops_per_gpu
1196 + or self.total_read_bytes_per_gpu
1197 + or self.total_write_bytes_per_gpu
1198 + ):
1199 + return
1200 +
1201 + now = time.monotonic()
1202 + delta_time = now - self.last_log_time
1203 +
1204 + if delta_time <= 0.0:
1205 + avg_tflops_per_gpu = 0.0
1206 + avg_gbps_per_gpu = 0.0
1207 + else:
1208 + avg_tflops_per_gpu = self.total_num_flops_per_gpu / delta_time / 1e12
1209 + avg_gbps_per_gpu = (
1210 + (self.total_read_bytes_per_gpu + self.total_write_bytes_per_gpu)
1211 + / delta_time
1212 + / 1e9
1213 + )
1214 +
1215 + log_fn(
1216 + "%sMFU: %.1f TF/s/GPU %.1f GB/s/GPU",
1217 + log_prefix,
1218 + avg_tflops_per_gpu,
1219 + avg_gbps_per_gpu,
1220 + )
1221 +
1222 + if self.debug_logging:
1223 + self.debug_logging.log(log_fn, log_prefix, delta_time)
1224 +
1225 + self.reset()
1226 +
1227 +
1228 +## util functions
1229 +
1230 +
1231 +def get_required(obj: object, attr: str):
1232 + """Get an attr from an object, or throw a InvalidComponentError if it's not set."""
1233 + if not hasattr(obj, attr):
1234 + raise InvalidComponent(f"Missing required attr {attr} in config")
1235 + return getattr(obj, attr)
1236 +
1237 +
1238 +def getattr_from_list(obj: object, attrs: list[str], default: object = None):
1239 + """Try to get the first attr that exists in the object
1240 + from a list of attrs. Otherwise return None."""
1241 + for attr in attrs:
1242 + if hasattr(obj, attr):
1243 + return getattr(obj, attr)
1244 + return default
vllm/v1/metrics/stats.py
+3
@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any
8
9 import vllm.envs as envs
10 from vllm.compilation.cuda_graph import CUDAGraphStat
11 +from vllm.v1.metrics.perf import PerfStats
12 from vllm.v1.spec_decode.metrics import SpecDecodingStats
13
14 if TYPE_CHECKING:
@@ -186,6 +187,8 @@ class SchedulerStats:
187
188 cudagraph_stats: CUDAGraphStat | None = None
189
190 + perf_stats: PerfStats | None = None
191 +
192
193 @dataclass
194 class RequestStateStats: