[test] Add LLM stats tests for pipeline parallel mode

- Introduced new test cases for LLM stats to validate behavior with multiple pipeline parallel configurations.
- Added micro batch ID tracking to LLM stats and verify it in the test cases.
- Used the new test cases to verify the new pipeline parallel mode behavior with chunked prefill enabled.

Signed-off-by: Robin Kobus <[email protected]>

Author: Funatiq

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -697,7 +697,7 @@ def get_queued_req_stats(request_id: int) -> RequestStats:
         return req_stats
 
     def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests,
-                           scheduled_batch) -> IterationStats:
+                           scheduled_batch, micro_batch_id) -> IterationStats:
         stats.iter_latency_ms = iter_latency_ms
 
         stats.num_queued_requests = self.executor_request_queue.get_request_queue_size(
@@ -738,7 +738,7 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests,
         stats.inflight_batching_stats.num_paused_requests = len(
             scheduled_batch.paused_requests)
         stats.inflight_batching_stats.avg_num_decoded_tokens_per_iter = 0
-        stats.inflight_batching_stats.micro_batch_id = 0
+        stats.inflight_batching_stats.micro_batch_id = micro_batch_id
         if stats.specdec_stats is not None:
             stats.specdec_stats.draft_overhead = 0.0 if iter_latency_ms <= 0.0 else float(
                 stats.specdec_stats.iter_latency_ms) / float(iter_latency_ms)
@@ -751,9 +751,13 @@ def _append_iter_stats(self,
         with self.stats_lock:
             self.stats.append((stats, req_stats))
 
-    def _process_iter_stats(self, finished_requests: list[LlmRequest],
-                            active_requests: List[LlmRequest],
-                            batch_state: BatchState):
+    def _process_iter_stats(
+        self,
+        finished_requests: list[LlmRequest],
+        active_requests: List[LlmRequest],
+        batch_state: BatchState,
+        micro_batch_id: int = 0,
+    ):
         iter_end_time = time.time()
         iter_latency_ms = (iter_end_time - batch_state.iter_start_time) * 1e3
         if batch_state.iter_stats is None:
@@ -766,9 +770,10 @@ def _process_iter_stats(self, finished_requests: list[LlmRequest],
                 and self.enable_iter_perf_stats) else None
 
         self._append_iter_stats(
-            self._update_iter_stats(
-                batch_state.iter_stats, iter_latency_ms, len(finished_requests),
-                batch_state.sample_state.scheduled_requests), req_stats)
+            self._update_iter_stats(batch_state.iter_stats, iter_latency_ms,
+                                    len(finished_requests),
+                                    batch_state.sample_state.scheduled_requests,
+                                    micro_batch_id), req_stats)
 
     def _executor_loop_cleanup(self):
 
@@ -828,6 +833,7 @@ def _executor_loop_pp(self):
                 self.num_scheduled_requests = scheduled_batch.batch_size
 
                 logger.debug(
+                    f'iteration {self.iter_counter}, microbatch {microbatch_id}, '
                     f'has {len(self.active_requests)} active_requests, '
                     f'scheduled {len(scheduled_batch.context_requests)} context requests and '
                     f'{len(scheduled_batch.generation_requests)} generation requests'
@@ -1008,9 +1014,11 @@ def _executor_loop_pp(self):
                 microbatch_id = (microbatch_id + 1) % self.num_micro_batches
 
                 if self.enable_iter_perf_stats and previous_batch is not None:
+                    sample_state = previous_batch.sample_state
+                    sample_state.scheduled_requests.context_requests = previous_batch.scheduled_ctx_reqs
                     self._process_iter_stats(finished_requests,
                                              self.active_requests,
-                                             previous_batch)
+                                             previous_batch, microbatch_id)
 
                 self.iter_counter += 1
 

tests/unittest/llmapi/test_llm.py

@@ -2050,59 +2050,76 @@ def validate_stats(
     results,
     pytorch_backend,
     max_tokens,
+    pp_size=1,
     use_overlap=False,
     enable_chunked_prefill=False,
     enable_iter_req_stats=False,
 ):
     assert results
+    for iter, result in enumerate(results):
+        ifbStats = result["inflightBatchingStats"]
+        print(f"iter: {iter}, ifbStats: {ifbStats}")
+
     expected_num_results = max_tokens if pytorch_backend else max_tokens + 1
     if enable_chunked_prefill:
         expected_num_results += 1
     assert len(results) == expected_num_results
 
     context_iterations = 2 if enable_chunked_prefill else 1
     generation_iterations = max_tokens - 1
+    microbatch_id = 0
     for iter, result in enumerate(results):
         ifbStats = result["inflightBatchingStats"]
 
         if iter < context_iterations:
-            assert ifbStats["numScheduledRequests"] == 1
-            assert ifbStats["numContextRequests"] == 1
-            assert ifbStats["numGenRequests"] == 0
-            assert result["numActiveRequests"] == 1
+            assert ifbStats["numScheduledRequests"] == 1, f"iter: {iter}"
+            assert ifbStats["numContextRequests"] == 1, f"iter: {iter}"
+            assert ifbStats["numGenRequests"] == 0, f"iter: {iter}"
+            assert result["numActiveRequests"] == 1, f"iter: {iter}"
+            assert ifbStats["microBatchId"] == microbatch_id, f"iter: {iter}"
         elif iter < (context_iterations + generation_iterations):
-            assert ifbStats["numScheduledRequests"] == 1
-            assert ifbStats["numContextRequests"] == 0
-            assert ifbStats["numGenRequests"] == 1
-            assert result["numActiveRequests"] == 1
+            assert ifbStats["numScheduledRequests"] == 1, f"iter: {iter}"
+            assert ifbStats["numContextRequests"] == 0, f"iter: {iter}"
+            assert ifbStats["numGenRequests"] == 1, f"iter: {iter}"
+            assert result["numActiveRequests"] == 1, f"iter: {iter}"
+            assert ifbStats["microBatchId"] == microbatch_id, f"iter: {iter}"
         else:
-            assert ifbStats["numScheduledRequests"] == 0
-            assert ifbStats["numContextRequests"] == 0
-            assert ifbStats["numGenRequests"] == 0
-            assert result["numActiveRequests"] == 0
+            assert ifbStats["numScheduledRequests"] == 0, f"iter: {iter}"
+            assert ifbStats["numContextRequests"] == 0, f"iter: {iter}"
+            assert ifbStats["numGenRequests"] == 0, f"iter: {iter}"
+            assert result["numActiveRequests"] == 0, f"iter: {iter}"
+            assert ifbStats["microBatchId"] == microbatch_id, f"iter: {iter}"
+
+        # In pipeline parallel mode, increment microbatch_id for each context iteration except the last one,
+        # since the context chunks can be scheduled in each iteration.
+        if pp_size > 1 and iter < context_iterations - 1:
+            microbatch_id += 1
 
         if enable_iter_req_stats:
-            assert "requestStats" in result
+            assert "requestStats" in result, f"iter: {iter}"
             req_stats = result["requestStats"]
-            assert len(req_stats) == 1
+            assert len(req_stats) == 1, f"iter: {iter}"
             req_stat = req_stats[0]
             if iter < (context_iterations - 1):
                 # If use_overlap, the stats are one iteration ahead
                 assert req_stat[
-                    "stage"] == "GENERATION_IN_PROGRESS" if use_overlap else "CONTEXT_IN_PROGRESS"
+                    "stage"] == "GENERATION_IN_PROGRESS" if use_overlap else "CONTEXT_IN_PROGRESS", f"iter: {iter}"
                 assert req_stat[
-                    "contextPrefillPosition"] == 54 if use_overlap else 32
-                assert req_stat["numGeneratedTokens"] == 0
+                    "contextPrefillPosition"] == 54 if use_overlap else 32, f"iter: {iter}"
+                assert req_stat["numGeneratedTokens"] == 0, f"iter: {iter}"
             elif iter < (context_iterations - 1 + generation_iterations):
-                assert req_stat["stage"] == "GENERATION_IN_PROGRESS"
-                assert req_stat["contextPrefillPosition"] == 54
+                assert req_stat[
+                    "stage"] == "GENERATION_IN_PROGRESS", f"iter: {iter}"
+                assert req_stat["contextPrefillPosition"] == 54, f"iter: {iter}"
                 assert req_stat["numGeneratedTokens"] == iter - (
-                    context_iterations - 1) + 1
+                    context_iterations - 1) + 1, f"iter: {iter}"
             else:
-                assert req_stat["stage"] == "GENERATION_COMPLETE"
-                assert req_stat["contextPrefillPosition"] == 54
-                assert req_stat["numGeneratedTokens"] == max_tokens
-            assert req_stat["scheduled"] == True
+                assert req_stat[
+                    "stage"] == "GENERATION_COMPLETE", f"iter: {iter}"
+                assert req_stat["contextPrefillPosition"] == 54, f"iter: {iter}"
+                assert req_stat[
+                    "numGeneratedTokens"] == max_tokens, f"iter: {iter}"
+            assert req_stat["scheduled"] == True, f"iter: {iter}"
 
         expected_num_completed = 1 if iter == len(results) - 1 else 0
 
@@ -2178,6 +2195,7 @@ def llm_get_stats_test_harness(tp_size: int = 1,
         results = llm.get_stats(2)
 
         validate_stats(results=results,
+                       pp_size=pp_size,
                        pytorch_backend=pytorch_backend,
                        max_tokens=max_tokens,
                        use_overlap=use_overlap,
@@ -2328,6 +2346,7 @@ async def task1():
             assert results
             if not use_overlap:
                 validate_stats(results=results,
+                               pp_size=pp_size,
                                pytorch_backend=pytorch_backend,
                                max_tokens=max_tokens,
                                use_overlap=use_overlap,

tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py

@@ -11,6 +11,7 @@
     check_llama_7b_multi_lora_from_request_test_harness,
     check_phi3_lora_fused_modules_output_tp2_identical_to_tp1)
 from .test_llm import (_test_llm_capture_request_error, llama_model_path,
+                       llm_get_stats_test_harness,
                        llm_return_logprobs_test_harness,
                        tinyllama_logits_processor_test_harness)
 from .test_llm_pytorch import llama_7b_lora_from_dir_test_harness
@@ -125,3 +126,45 @@ def test_llm_return_logprobs_streaming_tp2(prompt_logprobs, logprobs,
                                      streaming=True,
                                      backend="pytorch",
                                      tp_size=2)
+
+
+@skip_ray
+@pytest.mark.gpu2
+@pytest.mark.parametrize(
+    "return_context_logits, enable_chunked_prefill, enable_iter_req_stats",
+    [
+        (False, False, True),
+        (False, True, True),
+    ],
+)
+def test_llm_get_stats_pp2(return_context_logits, enable_chunked_prefill,
+                           enable_iter_req_stats):
+    llm_get_stats_test_harness(
+        tp_size=1,
+        pp_size=2,
+        return_context_logits=return_context_logits,
+        pytorch_backend=True,
+        enable_chunked_prefill=enable_chunked_prefill,
+        enable_iter_req_stats=enable_iter_req_stats,
+    )
+
+
+@skip_ray
+@pytest.mark.gpu4
+@pytest.mark.parametrize(
+    "return_context_logits, enable_chunked_prefill, enable_iter_req_stats",
+    [
+        (False, False, True),
+        (False, True, True),
+    ],
+)
+def test_llm_get_stats_pp4(return_context_logits, enable_chunked_prefill,
+                           enable_iter_req_stats):
+    llm_get_stats_test_harness(
+        tp_size=1,
+        pp_size=4,
+        return_context_logits=return_context_logits,
+        pytorch_backend=True,
+        enable_chunked_prefill=enable_chunked_prefill,
+        enable_iter_req_stats=enable_iter_req_stats,
+    )