Skip to content

Commit c17c423

Browse files
richardhuo-nvlancelly
authored andcommitted
[None][fix] Serialize the window_size in the kv event (NVIDIA#6526)
Signed-off-by: richardhuo-nv <[email protected]> Signed-off-by: Lanyu Liao <[email protected]>
1 parent 76513f0 commit c17c423

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
@@ -1016,6 +1016,7 @@ def to_json_str(cls, event):
10161016
return {
10171017
"event_id": event.event_id,
10181018
"data": event_serialize_func(event.data),
1019+
"window_size": event.window_size
10191020
}
10201021

10211022
@staticmethod

tests/unittest/llmapi/test_llm_kv_cache_events.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,8 @@ def test_kv_cache_event_data_serialization():
7171
flush_events(kv_cache_manager)
7272
events = kv_cache_manager.get_latest_events(10)
7373
serialized_event = KVCacheEventSerializer.serialize(events)
74-
assert len(serialized_event) == 1 and serialized_event[0]["event_id"] == 0
74+
assert len(serialized_event) == 1 and serialized_event[0][
75+
"event_id"] == 0 and serialized_event[0]["window_size"] == 256
7576
assert serialized_event[0]["data"]["type"] == "created"
7677
assert len(serialized_event[0]["data"]["num_blocks_per_cache_level"]) == 2
7778

0 commit comments

Comments
 (0)