Adding more tests

Signed-off-by: Patrice Castonguay <[email protected]>

Author: pcastonguay

cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp

@@ -125,7 +125,7 @@ void initConfigBindings(nb::module_& m)
             nb::cast<bool>(state[6]), nb::cast<std::optional<float>>(state[7]),
             nb::cast<std::optional<tle::RetentionPriority>>(state[8]), nb::cast<size_t>(state[9]),
             nb::cast<bool>(state[10]), nb::cast<bool>(state[11]), nb::cast<bool>(state[12]),
-            nb::cast<SizeTyp32>(state[13]));
+            nb::cast<SizeType32>(state[13]));
     };
     nb::class_<tle::KvCacheConfig>(m, "KvCacheConfig")
         .def(nb::init<bool, std::optional<SizeType32> const&, std::optional<std::vector<SizeType32>> const&,

cpp/tests/unit_tests/batch_manager/kvCacheEventManagerTest.cpp

@@ -915,6 +915,40 @@ void compareKvCacheEvents(texec::KVCacheEvent const& kvCacheEvent, texec::KVCach
     }
 }
 
+TEST(SerializeUtilsTest, KvCacheEventsDeque)
+{
+    // Created event
+    texec::KVCacheCreatedData kvCacheCreatedData{{1, 2}};
+    texec::KVCacheEvent kvCacheCreatedEvent(1, kvCacheCreatedData, 32);
+
+    // Removed event
+    texec::KVCacheEvent kvCacheRemovedEvent(1, texec::KVCacheRemovedData{{3, 4}}, 32);
+
+    // Stored event
+    auto storedBlockData1 = texec::KVCacheStoredBlockData(77, {{1, 2}, {3, 4}, {5, 6}}, 88, 0, 99);
+    auto storedBlockData2 = texec::KVCacheStoredBlockData(99, {{11, 12}, {3, 4}, {15, 6}}, 77, 1, 101);
+    texec::KVCacheStoredData kvCacheStoredData{177, {storedBlockData1, storedBlockData2}};
+    texec::KVCacheEvent kvCacheStoredEvent(1, kvCacheStoredData, 32);
+
+    // Updated event
+    texec::KVCacheEventDiff<texec::SizeType32> diff{0, 1};
+    texec::KVCacheEventDiff<texec::SizeType32> diff2{90, 99};
+    texec::KVCacheUpdatedData kvCacheUpdatedData(999, diff, diff2);
+    texec::KVCacheEvent kvCacheEvent(1, kvCacheUpdatedData, 32);
+
+    std::deque<texec::KVCacheEvent> kvCacheEvents{
+        kvCacheCreatedEvent, kvCacheRemovedEvent, kvCacheStoredEvent, kvCacheEvent};
+
+    auto serializedEvents = texec::Serialization::serialize(kvCacheEvents);
+    auto kvCacheEvents2 = texec::Serialization::deserializeKVCacheEvents(serializedEvents);
+
+    EXPECT_EQ(kvCacheEvents.size(), kvCacheEvents2.size());
+    for (size_t i = 0; i < kvCacheEvents.size(); ++i)
+    {
+        compareKvCacheEvents(kvCacheEvents[i], kvCacheEvents2[i]);
+    }
+}
+
 // Test for KVCacheEvent with KVCacheCreatedData
 TEST(SerializeUtilsTest, KVCacheCreatedEvent)
 {

cpp/tests/unit_tests/executor/serializeUtilsTest.cpp

@@ -301,15 +301,16 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
         }
         if self.event_buffer_max_size > 0:
             if mapping.enable_attention_dp:
-                 kwargs['event_manager'] = KVCacheEventManagerCpp(
-                     max_kv_event_entries=self.event_buffer_max_size,
-                     attention_dp_rank=mapping.rank,
-                     attention_dp_size=mapping.world_size,
-                     attention_dp_events_gather_period_ms=self.attention_dp_events_gather_period_ms,
-                 )
-             else:
-                 kwargs['event_manager'] = KVCacheEventManagerCpp(
-                     max_kv_event_entries=self.event_buffer_max_size)
+                kwargs['event_manager'] = KVCacheEventManagerCpp(
+                    max_kv_event_entries=self.event_buffer_max_size,
+                    attention_dp_rank=mapping.rank,
+                    attention_dp_size=mapping.world_size,
+                    attention_dp_events_gather_period_ms=self.
+                    attention_dp_events_gather_period_ms,
+                )
+            else:
+                kwargs['event_manager'] = KVCacheEventManagerCpp(
+                    max_kv_event_entries=self.event_buffer_max_size)
 
         self.impl = KVCacheManagerCpp(**kwargs)