[TRTLLM-11421][test] Add unit and integration tests for KV cache iteration stats

Unit tests:
- test_stats_serializer.py: 6 tests for kvCacheIterationStats serialization
  (field presence, multiple window sizes, missing stats, interval gating)
- test_collector.py: 5 tests for Prometheus metric collection
  (gauge updates, counter increments, multi-window labels, missing stats)
- _test_openai_metrics.py / _test_openai_prometheus.py: field additions for
  serving endpoint tests

Integration test (test_kv_cache_iteration_stats.py) with 8 pytest scenarios:
1. Cold start — verifies iterMissedBlocks, iterAllocTotalBlocks, iterGenAllocBlocks
2. Partial block reuse — short prompt repeated, verifies iterPartialReusedBlocks
3. Full block reuse — long prompt (3+ blocks) repeated, verifies iterFullReusedBlocks
4. Shared prefix — common prefix with different suffixes, verifies iterReusedBlocks
5. Batch generation — multiple prompts in one call, verifies pool usage
6. Long context — large prompt, verifies block allocation at scale
7. Rapid-fire — 20 requests, verifies accumulated deltas
8. Field completeness — all 18 fields present in every stats entry

Runnable via pytest or standalone. Supports --verbose-stats (pytest) or
--verbose/--test N/--list (standalone) for selective execution.

Added to premerge L0 pipeline (l0_h100.yml, l0_b200.yml).

Signed-off-by: Yueh-Ting Chen <yueh.ting.chen@gmail.com>

Author: eopXD

tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py

@@ -154,6 +154,10 @@ l0_b200:
   - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2LoRA::test_lora_multi_adapter_v2
   - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2LoRA::test_lora_chunked_prefill
   - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2LoRA::test_lora_eviction
+  # ------------- KV Cache Iteration Stats ---------------
+  - unittest/executor/test_stats_serializer.py
+  - unittest/metrics/test_collector.py
+  - kv_cache/test_kv_cache_iteration_stats.py
   # ------------- Visual Gen tests ---------------
   - unittest/_torch/visual_gen/test_visual_gen_args.py
   - unittest/_torch/visual_gen/test_teacache.py

tests/integration/test_lists/test-db/l0_b200.yml

@@ -428,6 +428,10 @@ l0_h100:
   - unittest/trt/attention/test_gpt_attention_no_cache.py
   - examples/test_gpt.py::test_gpt_oss_20b_lora_torch[gpt-oss-20b-lora-adapter_NIM_r8-gpt-oss-20b]
   - unittest/kv_cache_manager_v2_tests/ # 4 min
+  # ------------- KV Cache Iteration Stats ---------------
+  - unittest/executor/test_stats_serializer.py
+  - unittest/metrics/test_collector.py
+  - kv_cache/test_kv_cache_iteration_stats.py
 - condition:
     ranges:
       system_gpu_count:

tests/integration/test_lists/test-db/l0_h100.yml

@@ -0,0 +1,180 @@
+"""Tests for _stats_serializer with kvCacheIterationStats injection."""
+
+import json
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from tensorrt_llm.executor.base_worker import BaseWorker
+
+
+def _make_mock_iteration_stats(kv_cache_stats_json=None):
+    """Create a mock IterationStats object with to_json_str()."""
+    base = {
+        "iter": 1,
+        "iterLatencyMS": 10.5,
+        "gpuMemUsage": 1024,
+        "cpuMemUsage": 0,
+        "pinnedMemUsage": 0,
+    }
+    if kv_cache_stats_json is not None:
+        base["kvCacheStats"] = kv_cache_stats_json
+
+    mock = MagicMock()
+    mock.to_json_str.return_value = json.dumps(base)
+    return mock
+
+
+def _make_mock_kv_iter_stats(
+    window_size=16,
+    primary_used=10,
+    primary_max=20,
+    reused=5,
+    full_reused=4,
+    partial_reused=1,
+    missed=3,
+    gen_alloc=2,
+):
+    """Create a mock KvCacheIterationStats nanobind object."""
+    s = SimpleNamespace(
+        primary_max_num_blocks=primary_max,
+        primary_free_num_blocks=primary_max - primary_used,
+        primary_used_num_blocks=primary_used,
+        secondary_max_num_blocks=0,
+        secondary_free_num_blocks=0,
+        secondary_used_num_blocks=0,
+        iter_alloc_total_blocks=reused + missed,
+        iter_alloc_new_blocks=missed,
+        iter_reused_blocks=reused,
+        iter_full_reused_blocks=full_reused,
+        iter_partial_reused_blocks=partial_reused,
+        iter_missed_blocks=missed,
+        iter_cache_hit_rate=reused / (reused + missed) if (reused + missed) > 0 else 0.0,
+        iter_gen_alloc_blocks=gen_alloc,
+        iter_onboard_blocks=1,
+        iter_onboard_bytes=4096,
+        iter_offload_blocks=0,
+        iter_offload_bytes=0,
+    )
+    return {window_size: s}
+
+
+class TestStatsSerializer:
+    def test_serializer_without_kv_iter_stats(self):
+        """Legacy 2-tuple and 3-tuple with None should produce same output."""
+        iter_stats = _make_mock_iteration_stats()
+
+        # 3-tuple with None kv_iter_stats
+        result = BaseWorker._stats_serializer((iter_stats, None, None))
+        d = json.loads(result)
+        assert "iter" in d
+        assert "kvCacheIterationStats" not in d
+
+    def test_serializer_with_kv_iter_stats(self):
+        """KvCacheIterationStats should appear when provided."""
+        iter_stats = _make_mock_iteration_stats(
+            kv_cache_stats_json={"maxNumBlocks": 20, "usedNumBlocks": 10}
+        )
+        kv_iter = _make_mock_kv_iter_stats(
+            window_size=16,
+            primary_used=10,
+            primary_max=20,
+            reused=5,
+            full_reused=4,
+            partial_reused=1,
+            missed=3,
+            gen_alloc=2,
+        )
+
+        result = BaseWorker._stats_serializer((iter_stats, None, kv_iter))
+        d = json.loads(result)
+
+        # Existing kvCacheStats should still be present
+        assert "kvCacheStats" in d
+
+        # New kvCacheIterationStats should be present
+        assert "kvCacheIterationStats" in d
+        iter_kv = d["kvCacheIterationStats"]
+        assert "16" in iter_kv  # window size key as string
+
+        ws_stats = iter_kv["16"]
+        assert ws_stats["primaryMaxNumBlocks"] == 20
+        assert ws_stats["primaryUsedNumBlocks"] == 10
+        assert ws_stats["primaryFreeNumBlocks"] == 10
+        assert ws_stats["iterReusedBlocks"] == 5
+        assert ws_stats["iterFullReusedBlocks"] == 4
+        assert ws_stats["iterPartialReusedBlocks"] == 1
+        assert ws_stats["iterMissedBlocks"] == 3
+        assert ws_stats["iterGenAllocBlocks"] == 2
+        assert ws_stats["iterOnboardBlocks"] == 1
+        assert ws_stats["iterOnboardBytes"] == 4096
+        assert ws_stats["iterOffloadBlocks"] == 0
+        assert ws_stats["iterOffloadBytes"] == 0
+        assert ws_stats["iterCacheHitRate"] == pytest.approx(5 / 8)
+
+    def test_serializer_multiple_window_sizes(self):
+        """Multiple window sizes should all appear in output."""
+        iter_stats = _make_mock_iteration_stats()
+        kv_iter = _make_mock_kv_iter_stats(
+            window_size=16,
+            primary_used=5,
+            primary_max=10,
+            reused=2,
+            full_reused=2,
+            partial_reused=0,
+            missed=1,
+            gen_alloc=0,
+        )
+        # Add a second window size
+        kv_iter[64] = _make_mock_kv_iter_stats(
+            window_size=64,
+            primary_used=8,
+            primary_max=16,
+            reused=3,
+            full_reused=1,
+            partial_reused=2,
+            missed=2,
+            gen_alloc=1,
+        )[64]
+
+        result = BaseWorker._stats_serializer((iter_stats, None, kv_iter))
+        d = json.loads(result)
+
+        iter_kv = d["kvCacheIterationStats"]
+        assert "16" in iter_kv
+        assert "64" in iter_kv
+        assert iter_kv["16"]["primaryMaxNumBlocks"] == 10
+        assert iter_kv["64"]["primaryMaxNumBlocks"] == 16
+
+    def test_serializer_with_request_stats(self):
+        """Request stats and kv iter stats should coexist."""
+        iter_stats = _make_mock_iteration_stats()
+        kv_iter = _make_mock_kv_iter_stats()
+
+        req_stat = MagicMock()
+        req_stat.to_json_str.return_value = json.dumps({"id": 42})
+
+        result = BaseWorker._stats_serializer((iter_stats, [req_stat], kv_iter))
+        d = json.loads(result)
+
+        assert "requestStats" in d
+        assert len(d["requestStats"]) == 1
+        assert d["requestStats"][0]["id"] == 42
+        assert "kvCacheIterationStats" in d
+
+    def test_serializer_none_on_off_interval(self):
+        """When kv_iter_stats is None (off-interval), field should be absent."""
+        iter_stats = _make_mock_iteration_stats()
+
+        result = BaseWorker._stats_serializer((iter_stats, None, None))
+        d = json.loads(result)
+        assert "kvCacheIterationStats" not in d
+
+    def test_serializer_legacy_2_tuple(self):
+        """Legacy 2-tuple without third element should work."""
+        iter_stats = _make_mock_iteration_stats()
+
+        result = BaseWorker._stats_serializer((iter_stats, None))
+        d = json.loads(result)
+        assert "kvCacheIterationStats" not in d

tests/unittest/executor/test_stats_serializer.py

@@ -98,3 +98,19 @@ def test_metrics(client):
     assert "pinnedMemUsage" in response_dict
     assert "staticBatchingStats" in response_dict
     assert "timestamp" in response_dict
+    # Per-iteration KV cache stats (keyed by window size)
+    assert "kvCacheIterationStats" in response_dict
+    kv_iter = response_dict["kvCacheIterationStats"]
+    assert len(kv_iter) > 0
+    # Check fields in the first (and likely only) window size entry
+    ws_stats = next(iter(kv_iter.values()))
+    assert "primaryMaxNumBlocks" in ws_stats
+    assert "primaryUsedNumBlocks" in ws_stats
+    assert "iterReusedBlocks" in ws_stats
+    assert "iterFullReusedBlocks" in ws_stats
+    assert "iterPartialReusedBlocks" in ws_stats
+    assert "iterMissedBlocks" in ws_stats
+    assert "iterCacheHitRate" in ws_stats
+    assert "iterGenAllocBlocks" in ws_stats
+    assert "iterOnboardBlocks" in ws_stats
+    assert "iterOnboardBytes" in ws_stats

tests/unittest/llmapi/apps/_test_openai_metrics.py

@@ -119,3 +119,8 @@ def test_metrics_endpoint(server: RemoteOpenAIServer):
         f"Iteration stats metrics not found after waiting {max_wait_time}s"
     assert metric_prefix + "kv_cache_hit_rate" in data
     assert metric_prefix + "kv_cache_utilization" in data
+
+    # Per-iteration KV cache metrics
+    assert metric_prefix + "kv_cache_iter_reuse_rate" in data
+    assert metric_prefix + "kv_cache_missed_blocks_total" in data
+    assert metric_prefix + "kv_cache_gen_alloc_blocks_total" in data