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

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()