Address open-rabbit review comments

Signed-off-by: Simeng Liu <simengl@nvidia.com>

Author: SimengLiu-nv

cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp

@@ -351,7 +351,8 @@ std::tuple<RequestVector, RequestVector> GuaranteedNoEvictScheduler::impl(
                         crossSummary = crossKvCacheManager->analyzePrefixReuse(uniqueTokens, *req);
                     }
                 }
-                else if (isEncoderInit && crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse()
+                else if (mEnablePrefixAwareScheduling && isEncoderInit && crossKvCacheManager
+                    && crossKvCacheManager->isEnableBlockReuse()
                     && !crossKvCacheManager->getBlockManager().isVariableWindow())
                 {
                     // Encoder admission only needs the cross summary for reuse ordering.

cpp/tensorrt_llm/executor/schedulerConfig.cpp

@@ -20,6 +20,29 @@
 namespace tensorrt_llm::executor
 {
 
+namespace
+{
+
+bool dynamicBatchConfigsEqual(
+    std::optional<DynamicBatchConfig> const& lhs, std::optional<DynamicBatchConfig> const& rhs)
+{
+    if (lhs.has_value() != rhs.has_value())
+    {
+        return false;
+    }
+    if (!lhs.has_value())
+    {
+        return true;
+    }
+
+    return lhs->getEnableBatchSizeTuning() == rhs->getEnableBatchSizeTuning()
+        && lhs->getEnableMaxNumTokensTuning() == rhs->getEnableMaxNumTokensTuning()
+        && lhs->getDynamicBatchMovingAverageWindow() == rhs->getDynamicBatchMovingAverageWindow()
+        && lhs->getBatchSizeTable() == rhs->getBatchSizeTable();
+}
+
+} // namespace
+
 SchedulerConfig::SchedulerConfig(CapacitySchedulerPolicy capacitySchedulerPolicy,
     std::optional<ContextChunkingPolicy> contextChunkingPolicy, std::optional<DynamicBatchConfig> dynamicBatchConfig,
     bool enablePrefixAwareScheduling)
@@ -34,6 +57,7 @@ bool SchedulerConfig::operator==(SchedulerConfig const& other) const
 {
     return mCapacitySchedulerPolicy == other.mCapacitySchedulerPolicy
         && mContextChunkingPolicy == other.mContextChunkingPolicy
+        && dynamicBatchConfigsEqual(mDynamicBatchConfig, other.mDynamicBatchConfig)
         && mEnablePrefixAwareScheduling == other.mEnablePrefixAwareScheduling;
 }
 

cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp

@@ -1416,6 +1416,9 @@ TEST_F(CombinedSchedulerTest, PrefixAwareSchedulingDisabledKeepsReusableTokensZe
     EXPECT_TRUE(disaggInit1.empty());
     EXPECT_TRUE(paused1.empty());
     EXPECT_EQ(req1->getEstimatedReusableTokens(), 0);
+    auto const req1Scheduled
+        = std::any_of(scheduled1.begin(), scheduled1.end(), [](auto const& req) { return req->mRequestId == 1; });
+    ASSERT_TRUE(req1Scheduled) << "Capacity scheduler must admit req1 before micro batch budget filtering";
 
     RequestVector microBatchActive;
     for (auto& req : scheduled1)

cpp/tests/unit_tests/executor/serializeUtilsTest.cpp

@@ -505,13 +505,26 @@ TEST(SerializeUtilsTest, SchedulerConfig)
     texec::SchedulerConfig defaultSchedulerConfig;
     EXPECT_TRUE(defaultSchedulerConfig.getEnablePrefixAwareScheduling());
 
+    texec::DynamicBatchConfig dynamicBatchConfig{/*enableBatchSizeTuning=*/true,
+        /*enableMaxNumTokensTuning=*/true, /*dynamicBatchMovingAverageWindow=*/64};
     texec::SchedulerConfig schedulerConfig(texec::CapacitySchedulerPolicy::kMAX_UTILIZATION,
-        texec::ContextChunkingPolicy::kFIRST_COME_FIRST_SERVED, std::nullopt, false);
+        texec::ContextChunkingPolicy::kFIRST_COME_FIRST_SERVED, dynamicBatchConfig, false);
     auto schedulerConfig2 = serializeDeserialize(schedulerConfig);
     EXPECT_EQ(schedulerConfig.getCapacitySchedulerPolicy(), schedulerConfig2.getCapacitySchedulerPolicy());
     EXPECT_EQ(schedulerConfig.getContextChunkingPolicy(), schedulerConfig2.getContextChunkingPolicy());
-    EXPECT_FALSE(schedulerConfig2.getDynamicBatchConfig().has_value());
+    EXPECT_EQ(schedulerConfig, schedulerConfig2);
+    ASSERT_TRUE(schedulerConfig2.getDynamicBatchConfig().has_value());
+    EXPECT_TRUE(schedulerConfig2.getDynamicBatchConfig()->getEnableBatchSizeTuning());
+    EXPECT_TRUE(schedulerConfig2.getDynamicBatchConfig()->getEnableMaxNumTokensTuning());
+    EXPECT_EQ(schedulerConfig2.getDynamicBatchConfig()->getDynamicBatchMovingAverageWindow(), 64);
     EXPECT_FALSE(schedulerConfig2.getEnablePrefixAwareScheduling());
+
+    texec::SchedulerConfig differentDynamicBatchConfig(texec::CapacitySchedulerPolicy::kMAX_UTILIZATION,
+        texec::ContextChunkingPolicy::kFIRST_COME_FIRST_SERVED,
+        texec::DynamicBatchConfig{/*enableBatchSizeTuning=*/false, /*enableMaxNumTokensTuning=*/true,
+            /*dynamicBatchMovingAverageWindow=*/64},
+        false);
+    EXPECT_FALSE(schedulerConfig == differentDynamicBatchConfig);
 }
 
 TEST(SerializeUtilsTest, ParallelConfig)