[Fix] Update number of available pages after prefix cache free (#2409)

MasterJH5574 · web-flow · commit 97df6979dee5 · 2024-05-24T14:41:01.000-04:00
This PR fixes an issue that causes the inconsistency of CanPrefill
result from different models.
diff --git a/cpp/serve/engine_actions/batch_prefill_base.cc b/cpp/serve/engine_actions/batch_prefill_base.cc
@@ -109,6 +109,8 @@ BatchPrefillBaseActionObj::GetRequestStateEntriesToPrefill(EngineState estate) {
                              current_total_seq_len, num_running_rsentries, kv_state_kind,
                              sliding_window_enabled)) {
             if (!estate->prefix_cache->TryFreeMemory()) break;
+            // Update number of available pages after memory free.
+            num_available_pages = models_[i]->GetNumAvailablePages();
           }
           if (CanPrefill(estate, num_prefill_rsentries + 1 + num_child_to_activate,
                          total_input_length, total_required_pages, num_available_pages,

Original file line number	Diff line number	Diff line change
`@@ -109,6 +109,8 @@ BatchPrefillBaseActionObj::GetRequestStateEntriesToPrefill(EngineState estate) {`
`109`	`109`	`current_total_seq_len, num_running_rsentries, kv_state_kind,`
`110`	`110`	`sliding_window_enabled)) {`
`111`	`111`	`if (!estate->prefix_cache->TryFreeMemory()) break;`
	`112`	`+ // Update number of available pages after memory free.`
	`113`	`+ num_available_pages = models_[i]->GetNumAvailablePages();`
`112`	`114`	`}`
`113`	`115`	`if (CanPrefill(estate, num_prefill_rsentries + 1 + num_child_to_activate,`
`114`	`116`	`total_input_length, total_required_pages, num_available_pages,`