diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index e733a0331f6..1f37e4a3426 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -1016,6 +1016,7 @@ def to_json_str(cls, event): return { "event_id": event.event_id, "data": event_serialize_func(event.data), + "window_size": event.window_size } @staticmethod diff --git a/tests/unittest/llmapi/test_llm_kv_cache_events.py b/tests/unittest/llmapi/test_llm_kv_cache_events.py index f5efbe2bcf8..8f7fb75c7f4 100644 --- a/tests/unittest/llmapi/test_llm_kv_cache_events.py +++ b/tests/unittest/llmapi/test_llm_kv_cache_events.py @@ -71,7 +71,8 @@ def test_kv_cache_event_data_serialization(): flush_events(kv_cache_manager) events = kv_cache_manager.get_latest_events(10) serialized_event = KVCacheEventSerializer.serialize(events) - assert len(serialized_event) == 1 and serialized_event[0]["event_id"] == 0 + assert len(serialized_event) == 1 and serialized_event[0][ + "event_id"] == 0 and serialized_event[0]["window_size"] == 256 assert serialized_event[0]["data"]["type"] == "created" assert len(serialized_event[0]["data"]["num_blocks_per_cache_level"]) == 2