Skip to content

Commit 1f35d90

Browse files
richardhuo-nvsymphonylyh
authored andcommitted
[None][fix] Serialize the window_size in the kv event (NVIDIA#6526)
Signed-off-by: richardhuo-nv <[email protected]>
1 parent 7b7059e commit 1f35d90

File tree

2 files changed

+3
-1
lines changed

2 files changed

+3
-1
lines changed

tensorrt_llm/_utils.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -991,6 +991,7 @@ def to_json_str(cls, event):
991991
return {
992992
"event_id": event.event_id,
993993
"data": event_serialize_func(event.data),
994+
"window_size": event.window_size
994995
}
995996

996997
@staticmethod

tests/unittest/llmapi/test_llm_kv_cache_events.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,8 @@ def test_kv_cache_event_data_serialization():
7575
flush_events(kv_cache_manager)
7676
events = kv_cache_manager.get_latest_events(10)
7777
serialized_event = KVCacheEventSerializer.serialize(events)
78-
assert len(serialized_event) == 1 and serialized_event[0]["event_id"] == 0
78+
assert len(serialized_event) == 1 and serialized_event[0][
79+
"event_id"] == 0 and serialized_event[0]["window_size"] == 256
7980
assert serialized_event[0]["data"]["type"] == "created"
8081
assert len(serialized_event[0]["data"]["num_blocks_per_cache_level"]) == 2
8182

0 commit comments

Comments
 (0)