From 4526a65a7f53d52d73924bee080aca8baa18f375 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:11:46 -0700 Subject: [PATCH 1/3] [None][fix] Enforce the V2 KV cache iteration stats contract Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../kv_cache/test_kv_cache_iteration_stats.py | 91 +++++++++++++------ .../executor/test_stats_serializer.py | 3 + 2 files changed, 64 insertions(+), 30 deletions(-) diff --git a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py b/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py index cb3826957661..a6865e6d8f81 100644 --- a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py +++ b/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py @@ -12,9 +12,9 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Integration tests for per-iteration KV cache statistics (kvCacheIterationStats). +"""Integration tests for KV cache manager V2 per-iteration statistics. -Tests verify that the 18 stat fields are correctly populated across +Tests verify that window, hot-pool and cold-pool fields are correctly populated across different inference scenarios: cold start, block reuse (partial/full), shared prefix, batch generation, long context, and rapid-fire. @@ -43,16 +43,13 @@ from ..conftest import llm_models_root MODEL = f"{llm_models_root()}/llama-models-v2/TinyLlama-1.1B-Chat-v1.0" +HOST_CACHE_SIZE = 64 << 20 -ALL_FIELDS = [ +WINDOW_FIELDS = { # Instantaneous gauges — primary (GPU) pool "primaryMaxNumBlocks", "primaryFreeNumBlocks", "primaryUsedNumBlocks", - # Instantaneous gauges — secondary (host) pool - "secondaryMaxNumBlocks", - "secondaryFreeNumBlocks", - "secondaryUsedNumBlocks", # Per-iteration deltas — context phase "iterAllocTotalBlocks", "iterAllocNewBlocks", @@ -71,14 +68,22 @@ # Intra-device (GPU → GPU) block copies "iterIntraDeviceCopyBlocks", "iterIntraDeviceCopyBytes", -] +} + +# Reuse counters describe logical cache lifecycles, not physical hot pools. +HOT_POOL_FIELDS = WINDOW_FIELDS - { + "iterReusedBlocks", + "iterFullReusedBlocks", + "iterPartialReusedBlocks", + "iterMissedBlocks", + "iterCacheHitRate", +} SECONDARY_FIELDS = { "secondaryMaxNumBlocks", "secondaryFreeNumBlocks", "secondaryUsedNumBlocks", } -NON_SECONDARY_FIELDS = set(ALL_FIELDS) - SECONDARY_FIELDS TEST_NAMES = { 1: "Cold start", @@ -103,21 +108,23 @@ def _is_verbose(request): def print_kv_stats(label, stats_list): - """Print all 18 fields for every stats entry.""" + """Print the fields reported in each V2 stats view.""" print(f"\n{'=' * 60}") print(f" {label}: {len(stats_list)} stats entries") print(f"{'=' * 60}") found = False for i, s in enumerate(stats_list): - ki = s.get("kvCacheIterationStats") - if ki: - found = True - for ws, v in ki.items(): - print(f"\n --- entry[{i}] window_size={ws} ---") - for field in ALL_FIELDS: - val = v.get(field, "") - print(f" {field:30s} = {val}") - else: + for view in ( + "kvCacheIterationStats", + "kvCacheIterationStatsByPoolGroup", + "kvCacheIterationStatsByColdPoolGroup", + ): + for group, fields in s.get(view, {}).items(): + found = True + print(f"\n --- entry[{i}] {view}[{group}] ---") + for field, value in fields.items(): + print(f" {field:30s} = {value}") + if not s.get("kvCacheIterationStats"): keys = list(s.keys())[:8] print(f" entry[{i}]: no kvCacheIterationStats (keys: {keys})") if not found: @@ -150,7 +157,12 @@ def llm_instance(): """Create a shared LLM instance for all tests in this module.""" llm = LLM( model=MODEL, - kv_cache_config=KvCacheConfig(enable_block_reuse=True, iteration_stats_interval=1), + kv_cache_config=KvCacheConfig( + version=2, + enable_block_reuse=True, + iteration_stats_interval=1, + host_cache_size=HOST_CACHE_SIZE, + ), enable_iter_perf_stats=True, return_perf_metrics=True, ) @@ -390,7 +402,7 @@ def test_rapid_fire(self, llm_instance, all_collected, request): assert total_alloc > 0, "iterAllocTotalBlocks = 0 across all entries" def test_field_completeness(self, llm_instance, all_collected, request): - """Field completeness — verify fields in their V2 window and cold-pool views.""" + """Verify fields in the V2 window, hot-pool and cold-pool views.""" # If running standalone (no prior tests), generate some traffic if not all_collected: llm_instance.generate(["Hello world"], SamplingParams(max_tokens=16)) @@ -401,13 +413,22 @@ def test_field_completeness(self, llm_instance, all_collected, request): ki = s.get("kvCacheIterationStats") if ki: entries_with_kv += 1 + assert s.get("kvCacheIterationStatsByPoolGroup"), "missing V2 hot-pool stats" + # The explicit host tier ensures cold-field coverage is not vacuous. + assert s.get("kvCacheIterationStatsByColdPoolGroup"), "missing V2 cold-pool stats" # V2 reports secondary gauges by cold pool group, not by window. - for ws, v in ki.items(): - missing_fields = NON_SECONDARY_FIELDS - v.keys() - assert not missing_fields, ( - f"Missing kvCacheIterationStats fields for window {ws}: " - f"{sorted(missing_fields)}" - ) + for view, expected_fields in ( + ("kvCacheIterationStats", WINDOW_FIELDS), + ("kvCacheIterationStatsByPoolGroup", HOT_POOL_FIELDS), + ): + for group, v in s[view].items(): + missing_fields = expected_fields - v.keys() + assert not missing_fields, ( + f"Missing {view} fields for group {group}: {sorted(missing_fields)}" + ) + assert not any(field.startswith("secondary") for field in v), ( + f"Unexpected secondary fields in {view} for group {group}" + ) for group, v in s.get("kvCacheIterationStatsByColdPoolGroup", {}).items(): missing_fields = SECONDARY_FIELDS - v.keys() @@ -415,6 +436,11 @@ def test_field_completeness(self, llm_instance, all_collected, request): f"Missing kvCacheIterationStatsByColdPoolGroup fields for group {group}: " f"{sorted(missing_fields)}" ) + assert v["secondaryMaxNumBlocks"] > 0 + assert ( + v["secondaryFreeNumBlocks"] + v["secondaryUsedNumBlocks"] + == v["secondaryMaxNumBlocks"] + ) print(f" Entries with kvCacheIterationStats: {entries_with_kv}/{len(all_collected)}") assert entries_with_kv > 0, "no entries contain kvCacheIterationStats" @@ -444,7 +470,7 @@ def main(): "--verbose", "-v", action="store_true", - help="Dump all 18 KV cache stat fields for every stats entry", + help="Dump the window, hot-pool and cold-pool KV cache stat fields", ) parser.add_argument( "--test", @@ -478,10 +504,15 @@ class FakeRequest: fake_request = FakeRequest() - print("Starting LLM with block_reuse + iteration_stats_interval=1") + print("Starting LLM with KV cache V2 + block_reuse + iteration_stats_interval=1") llm = LLM( model=MODEL, - kv_cache_config=KvCacheConfig(enable_block_reuse=True, iteration_stats_interval=1), + kv_cache_config=KvCacheConfig( + version=2, + enable_block_reuse=True, + iteration_stats_interval=1, + host_cache_size=HOST_CACHE_SIZE, + ), enable_iter_perf_stats=True, return_perf_metrics=True, ) diff --git a/tests/unittest/executor/test_stats_serializer.py b/tests/unittest/executor/test_stats_serializer.py index 45b31e7ea710..2d6418f298e6 100644 --- a/tests/unittest/executor/test_stats_serializer.py +++ b/tests/unittest/executor/test_stats_serializer.py @@ -448,6 +448,9 @@ def test_serializer_with_v2_pool_group_stats(self): assert cold_group["coldPoolGroupId"] == 0 assert cold_group["slotSize"] == [4 << 20] assert cold_group["windowSizes"] == [16, 64] + assert cold_group["secondaryMaxNumBlocks"] == 8 + assert cold_group["secondaryFreeNumBlocks"] == 5 + assert cold_group["secondaryUsedNumBlocks"] == 3 assert cold_group["secondaryPeakFreeNumBlocks"] == 6 assert cold_group["secondaryPeakUsedNumBlocks"] == 4 assert cold_group["secondaryPeakEvictableNumBlocks"] == 2 From a2a4d690ed16c1f09ace8d45a341ad1f6bca5cf7 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:53:07 -0700 Subject: [PATCH 2/3] [None][fix] Select V2 explicitly and narrow stats coverage Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../kv_cache/test_kv_cache_iteration_stats.py | 79 ++++++++----------- tests/unittest/metrics/test_collector.py | 8 +- 2 files changed, 39 insertions(+), 48 deletions(-) diff --git a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py b/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py index a6865e6d8f81..1536680cd574 100644 --- a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py +++ b/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py @@ -12,9 +12,9 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Integration tests for KV cache manager V2 per-iteration statistics. +"""Integration tests for per-iteration KV cache statistics (kvCacheIterationStats). -Tests verify that window, hot-pool and cold-pool fields are correctly populated across +Tests verify that the 18 stat fields are correctly populated across different inference scenarios: cold start, block reuse (partial/full), shared prefix, batch generation, long context, and rapid-fire. @@ -43,13 +43,18 @@ from ..conftest import llm_models_root MODEL = f"{llm_models_root()}/llama-models-v2/TinyLlama-1.1B-Chat-v1.0" +# Provision a host tier so the cold-pool field checks cannot pass vacuously. HOST_CACHE_SIZE = 64 << 20 -WINDOW_FIELDS = { +ALL_FIELDS = [ # Instantaneous gauges — primary (GPU) pool "primaryMaxNumBlocks", "primaryFreeNumBlocks", "primaryUsedNumBlocks", + # Instantaneous gauges — secondary (host) pool + "secondaryMaxNumBlocks", + "secondaryFreeNumBlocks", + "secondaryUsedNumBlocks", # Per-iteration deltas — context phase "iterAllocTotalBlocks", "iterAllocNewBlocks", @@ -68,22 +73,14 @@ # Intra-device (GPU → GPU) block copies "iterIntraDeviceCopyBlocks", "iterIntraDeviceCopyBytes", -} - -# Reuse counters describe logical cache lifecycles, not physical hot pools. -HOT_POOL_FIELDS = WINDOW_FIELDS - { - "iterReusedBlocks", - "iterFullReusedBlocks", - "iterPartialReusedBlocks", - "iterMissedBlocks", - "iterCacheHitRate", -} +] SECONDARY_FIELDS = { "secondaryMaxNumBlocks", "secondaryFreeNumBlocks", "secondaryUsedNumBlocks", } +NON_SECONDARY_FIELDS = set(ALL_FIELDS) - SECONDARY_FIELDS TEST_NAMES = { 1: "Cold start", @@ -108,23 +105,21 @@ def _is_verbose(request): def print_kv_stats(label, stats_list): - """Print the fields reported in each V2 stats view.""" + """Print all 18 fields for every stats entry.""" print(f"\n{'=' * 60}") print(f" {label}: {len(stats_list)} stats entries") print(f"{'=' * 60}") found = False for i, s in enumerate(stats_list): - for view in ( - "kvCacheIterationStats", - "kvCacheIterationStatsByPoolGroup", - "kvCacheIterationStatsByColdPoolGroup", - ): - for group, fields in s.get(view, {}).items(): - found = True - print(f"\n --- entry[{i}] {view}[{group}] ---") - for field, value in fields.items(): - print(f" {field:30s} = {value}") - if not s.get("kvCacheIterationStats"): + ki = s.get("kvCacheIterationStats") + if ki: + found = True + for ws, v in ki.items(): + print(f"\n --- entry[{i}] window_size={ws} ---") + for field in ALL_FIELDS: + val = v.get(field, "") + print(f" {field:30s} = {val}") + else: keys = list(s.keys())[:8] print(f" entry[{i}]: no kvCacheIterationStats (keys: {keys})") if not found: @@ -158,7 +153,7 @@ def llm_instance(): llm = LLM( model=MODEL, kv_cache_config=KvCacheConfig( - version=2, + use_kv_cache_manager_v2=True, enable_block_reuse=True, iteration_stats_interval=1, host_cache_size=HOST_CACHE_SIZE, @@ -402,7 +397,7 @@ def test_rapid_fire(self, llm_instance, all_collected, request): assert total_alloc > 0, "iterAllocTotalBlocks = 0 across all entries" def test_field_completeness(self, llm_instance, all_collected, request): - """Verify fields in the V2 window, hot-pool and cold-pool views.""" + """Field completeness — verify fields in their V2 window and cold-pool views.""" # If running standalone (no prior tests), generate some traffic if not all_collected: llm_instance.generate(["Hello world"], SamplingParams(max_tokens=16)) @@ -416,19 +411,15 @@ def test_field_completeness(self, llm_instance, all_collected, request): assert s.get("kvCacheIterationStatsByPoolGroup"), "missing V2 hot-pool stats" # The explicit host tier ensures cold-field coverage is not vacuous. assert s.get("kvCacheIterationStatsByColdPoolGroup"), "missing V2 cold-pool stats" + if entries_with_kv == 1: + print(f" V2 cold-pool stats: {s['kvCacheIterationStatsByColdPoolGroup']}") # V2 reports secondary gauges by cold pool group, not by window. - for view, expected_fields in ( - ("kvCacheIterationStats", WINDOW_FIELDS), - ("kvCacheIterationStatsByPoolGroup", HOT_POOL_FIELDS), - ): - for group, v in s[view].items(): - missing_fields = expected_fields - v.keys() - assert not missing_fields, ( - f"Missing {view} fields for group {group}: {sorted(missing_fields)}" - ) - assert not any(field.startswith("secondary") for field in v), ( - f"Unexpected secondary fields in {view} for group {group}" - ) + for ws, v in ki.items(): + missing_fields = NON_SECONDARY_FIELDS - v.keys() + assert not missing_fields, ( + f"Missing kvCacheIterationStats fields for window {ws}: " + f"{sorted(missing_fields)}" + ) for group, v in s.get("kvCacheIterationStatsByColdPoolGroup", {}).items(): missing_fields = SECONDARY_FIELDS - v.keys() @@ -437,10 +428,6 @@ def test_field_completeness(self, llm_instance, all_collected, request): f"{sorted(missing_fields)}" ) assert v["secondaryMaxNumBlocks"] > 0 - assert ( - v["secondaryFreeNumBlocks"] + v["secondaryUsedNumBlocks"] - == v["secondaryMaxNumBlocks"] - ) print(f" Entries with kvCacheIterationStats: {entries_with_kv}/{len(all_collected)}") assert entries_with_kv > 0, "no entries contain kvCacheIterationStats" @@ -470,7 +457,7 @@ def main(): "--verbose", "-v", action="store_true", - help="Dump the window, hot-pool and cold-pool KV cache stat fields", + help="Dump all 18 KV cache stat fields for every stats entry", ) parser.add_argument( "--test", @@ -504,11 +491,11 @@ class FakeRequest: fake_request = FakeRequest() - print("Starting LLM with KV cache V2 + block_reuse + iteration_stats_interval=1") + print("Starting LLM with block_reuse + iteration_stats_interval=1") llm = LLM( model=MODEL, kv_cache_config=KvCacheConfig( - version=2, + use_kv_cache_manager_v2=True, enable_block_reuse=True, iteration_stats_interval=1, host_cache_size=HOST_CACHE_SIZE, diff --git a/tests/unittest/metrics/test_collector.py b/tests/unittest/metrics/test_collector.py index 7a9558ba3864..f3fe88d3ab27 100644 --- a/tests/unittest/metrics/test_collector.py +++ b/tests/unittest/metrics/test_collector.py @@ -833,14 +833,18 @@ def test_v2_lifecycle_and_pool_group_stats_are_aggregated(self): }, "kvCacheIterationStatsByPoolGroup": { "0": { - "secondaryMaxNumBlocks": 50, - "secondaryUsedNumBlocks": 20, "iterGenAllocBlocks": 2, "iterOnboardBytes": 4096, "iterOffloadBytes": 2048, "iterIntraDeviceCopyBytes": 8192, } }, + "kvCacheIterationStatsByColdPoolGroup": { + "0": { + "secondaryMaxNumBlocks": 50, + "secondaryUsedNumBlocks": 20, + } + }, } before_reused = _get_counter_value(collector, "kv_cache_iter_reused_blocks") From 8614aa729eae917f734a99170943a84b15e67740 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 15 Sep 2026 19:07:07 -0700 Subject: [PATCH 3/3] [None][test] Clarify host-tier coverage and cold-pool diagnostics Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../defs/kv_cache/test_kv_cache_iteration_stats.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py b/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py index 1536680cd574..3a7d62c85860 100644 --- a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py +++ b/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py @@ -43,7 +43,7 @@ from ..conftest import llm_models_root MODEL = f"{llm_models_root()}/llama-models-v2/TinyLlama-1.1B-Chat-v1.0" -# Provision a host tier so the cold-pool field checks cannot pass vacuously. +# Pin the host tier for all scenarios so cold-pool field coverage is deterministic. HOST_CACHE_SIZE = 64 << 20 ALL_FIELDS = [ @@ -149,7 +149,7 @@ def find_kv_entries(stats_list): # --------------------------------------------------------------------------- @pytest.fixture(scope="module") def llm_instance(): - """Create a shared LLM instance for all tests in this module.""" + """Share one V2 LLM with a fixed host tier across all eight scenarios.""" llm = LLM( model=MODEL, kv_cache_config=KvCacheConfig( @@ -427,7 +427,9 @@ def test_field_completeness(self, llm_instance, all_collected, request): f"Missing kvCacheIterationStatsByColdPoolGroup fields for group {group}: " f"{sorted(missing_fields)}" ) - assert v["secondaryMaxNumBlocks"] > 0 + assert v["secondaryMaxNumBlocks"] > 0, ( + f"Cold group {group} reports no secondary capacity: {v}" + ) print(f" Entries with kvCacheIterationStats: {entries_with_kv}/{len(all_collected)}") assert entries_with_kv > 0, "no entries contain kvCacheIterationStats"