[metrics] Fix prefill FLOPs estimate to count prefix and per-request causal pairs (#34316)

python/sglang/srt/managers/scheduler_components/metrics_reporter.py

@@ -459,6 +459,14 @@ class SchedulerMetricsReporter:
             num_attn_heads * head_dim * act_bytes * num_layers
         )
 
+    @staticmethod
+    def _prefill_attention_pairs(batch) -> float:
+        """Causal query-key pairs: each chunk against its cached prefix, plus
+        the causal pairs within the chunk itself."""
+        prefix_pairs = sum(c * p for c, p in zip(batch.extend_lens, batch.prefix_lens))
+        within_chunk_pairs = sum(c * (c + 1) / 2.0 for c in batch.extend_lens)
+        return float(prefix_pairs + within_chunk_pairs)
+
     def _estimate_prefill_perf(self, batch) -> Tuple[float, float, float]:
         if batch is None or batch.extend_lens is None:
             return 0.0, 0.0, 0.0
@@ -466,17 +474,20 @@ class SchedulerMetricsReporter:
         if tokens == 0:
             return 0.0, 0.0, 0.0
 
-        # Causal prefill token-context product.
-        context_product = tokens * (tokens + 1) / 2.0
+        context_product = self._prefill_attention_pairs(batch)
         flops = (
             tokens * self._linear_flops_per_token
             + self._attn_dot_flops_coeff * context_product
         )
 
+        # The chunk's queries share one pass over the cached prefix, so charge the
+        # prefix KV once per chunk -- not once per query-key pair.
+        prefix_kv_tokens = float(sum(batch.prefix_lens))
         read_bytes = (
             tokens * self._weight_read_bytes_per_token
             + tokens * self._qkv_act_bytes_per_token
             + tokens * self._prefill_attn_act_read_per_token
+            + prefix_kv_tokens * self._kv_cache_bytes_per_token
         )
         write_bytes = (
             tokens * self._kv_cache_bytes_per_token
@@ -512,8 +523,8 @@ class SchedulerMetricsReporter:
 
     def _prefill_sol_suffix(self, batch, elapsed_s: float) -> str:
         """Hook: model-specific speed-of-light % suffix for the prefill log line.
-        ``batch`` carries the per-request extend/prefix lengths a subclass needs
-        for an exact attention pair-count. No model arch here, so returns "";
+        Call ``_prefill_attention_pairs(batch)`` for the exact causal
+        attention pair-count. No model arch here, so returns "";
         a subclass may override it."""
         return ""
 

test/registered/unit/observability/test_forward_pass_metrics.py

@@ -14,6 +14,7 @@ from sglang.srt.managers.scheduler_components.metrics_reporter import (
     PrefillStats,
     SchedulerMetricsReporter,
 )
+from sglang.test.test_utils import CustomTestCase
 
 
 def _make_ps(**overrides) -> ParallelState:
@@ -399,5 +400,51 @@ class TestIdleMetrics(unittest.TestCase):
         self.assertEqual(self.published_occupancies, [])
 
 
+class TestEstimatedPrefillPerf(CustomTestCase):
+    """Causal pair count behind ``est. prefill TFLOPS/s`` and ``estimated_flops``."""
+
+    def setUp(self):
+        self.scheduler = types.SimpleNamespace()
+        self.scheduler.waiting_queue = []
+        self.scheduler.disaggregation_mode = DisaggregationMode.NULL
+        self.reporter = _make_reporter(self, self.scheduler)
+        # One unit per query-key pair and nothing else, so the returned FLOPs
+        # are exactly the attention pair count.
+        self.reporter._linear_flops_per_token = 0.0
+        self.reporter._attn_dot_flops_coeff = 1.0
+        self.reporter._weight_read_bytes_per_token = 0.0
+        self.reporter._qkv_act_bytes_per_token = 0.0
+        self.reporter._prefill_attn_act_read_per_token = 0.0
+        self.reporter._kv_cache_bytes_per_token = 0.0
+        self.reporter._ffn_act_bytes_per_token = 0.0
+
+    def _pair_count(self, extend_lens, prefix_lens):
+        batch = types.SimpleNamespace(extend_lens=extend_lens, prefix_lens=prefix_lens)
+        flops, _, _ = self.reporter._estimate_prefill_perf(batch)
+        return flops
+
+    def test_chunk_is_charged_for_its_cached_prefix(self):
+        self.assertEqual(self._pair_count([4], [3]), 4 * 3 + 4 * 5 / 2)
+
+    def test_prefix_kv_is_read_once_per_chunk(self):
+        # One pass over the prefix per chunk, not one read per query-key pair:
+        # the chunk's queries share the same KV stream.
+        self.reporter._kv_cache_bytes_per_token = 1.0
+        batch = types.SimpleNamespace(extend_lens=[4], prefix_lens=[3])
+        _, read_bytes, _ = self.reporter._estimate_prefill_perf(batch)
+        self.assertEqual(read_bytes, 3)
+
+    def test_requests_in_one_batch_do_not_attend_to_each_other(self):
+        self.assertEqual(self._pair_count([100, 100], [0, 0]), 2 * (100 * 101 / 2))
+
+    def test_mixed_prefill_and_decode_rows_use_their_own_context(self):
+        # mix_with_running appends running requests as extend_len 1 with their
+        # full context as prefix_len.
+        self.assertEqual(
+            self._pair_count([8, 1, 1], [0, 100, 200]),
+            8 * 9 / 2 + (100 + 1) + (200 + 1),
+        )
+
+
 if __name__ == "__main__":
     unittest.main()