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 9, 13:33 UTC
[Metrics][MFU] Fix UnembedMetrics FLOP overcounting for prefill (#33045) (#33045)
Fix UnembedMetrics to correctly count FLOPs for the unembedding (LM head) layer. The bug: UnembedMetrics used total_num_tokens() which counts all tokens in the batch for projection flops, vocab projections are run on just the last token for the autoregressive use case. Co-authored-by: Omar Mohamed Khalil <omarkhalil@meta.com>
omkhalil committed
Jan 27, 2026 at 10:16 UTC
5ec44056f7f1ce81a8cd089249449f2fd4256fdb
1 file changed
+11
-3
vllm/v1/metrics/perf.py
+11
-3
@@ -110,6 +110,14 @@ class ExecutionContext:
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
+ def num_logits_tokens(self) -> int:
114
+ """Number of tokens that require logits computation (unembedding).
115
+
116
+ For prefill, only the last token per request needs logits.
117
+ For decode, all tokens need logits.
118
+ """
119
+ return self.num_prefill_requests + self.decode_num_tokens
120
+
121
@classmethod
122
def from_single_request(
123
cls, num_tokens: int, context_len: int, is_prefill: bool
@@ -906,7 +914,7 @@ class UnembedMetrics(ComponentMetrics):
914
) -> dict[str, int]:
915
"""Calculate flops breakdown for unembedding layer."""
916
D, V = self.hidden_size, self.vocab_size
909
- T = ctx.total_num_tokens()
917
+ T = ctx.num_logits_tokens()
918
919
if per_gpu:
920
V //= self.tp_size
@@ -920,7 +928,7 @@ class UnembedMetrics(ComponentMetrics):
928
) -> dict[str, int]:
929
"""Calculate read memory traffic for unembedding layer."""
930
D, V = self.hidden_size, self.vocab_size
923
- T = ctx.total_num_tokens()
931
+ T = ctx.num_logits_tokens()
932
933
if per_gpu:
934
V //= self.tp_size
@@ -935,7 +943,7 @@ class UnembedMetrics(ComponentMetrics):
943
) -> dict[str, int]:
944
"""Calculate write memory traffic for unembedding layer."""
945
V = self.vocab_size
938
- T = ctx.total_num_tokens()
946
+ T = ctx.num_logits_tokens()
947
948
if per_gpu:
949
V //= self.tp_size