Enable prefix chunking though the config overlay

Author: anujj

examples/python/model-generate.py

@@ -19,7 +19,8 @@ def main(args):
     batch_size = len(prompts)
 
     config = og.Config(args.model_path)
-    config.overlay(f'{{"search": {{"batch_size": {batch_size}, "num_beams": {3}}}}}')
+    # Example: Configure search parameters including chunk_size for prefix chunking  
+    config.overlay(f'{{"search": {{"batch_size": {batch_size}, "num_beams": {3}, "chunk_size": {args.chunk_size}}}}}')
 
     if args.execution_provider != "follow_config":
         config.clear_providers()
@@ -90,6 +91,7 @@ def main(args):
     parser.add_argument('-v', '--verbose', action='store_true', default=False, help='Print verbose output and timing information. Defaults to false')
     parser.add_argument('-b', '--batch_size_for_cuda_graph', type=int, default=1, help='Max batch size for CUDA graph')
     parser.add_argument('-c', '--chat_template', type=str, default='', help='Chat template to use for the prompt. User input will be injected into {input}. If not set, the prompt is used as is.')
+    parser.add_argument('--chunk_size', type=int, default=-1, help='Chunk size for prefix chunking during context processing (default: -1 = disabled, >0 = enabled)')
     parser.add_argument('--non-interactive', action=argparse.BooleanOptionalAction, required=False, default=False, help='Non-interactive mode, mainly for CI usage')
 
     args = parser.parse_args()

src/config.cpp

@@ -778,6 +778,8 @@ struct Search_Element : JSON::Element {
       v_.length_penalty = static_cast<float>(JSON::Get<double>(value));
     } else if (name == "random_seed") {
       v_.random_seed = static_cast<int>(JSON::Get<double>(value));
+    } else if (name == "chunk_size") {
+      v_.chunk_size = static_cast<int>(JSON::Get<double>(value));
     } else if (name == "do_sample") {
       v_.do_sample = JSON::Get<bool>(value);
     } else if (name == "past_present_share_buffer") {

src/config.h

@@ -266,6 +266,7 @@ struct Config {
     float length_penalty{1.0f};        // Exponential penalty to the length that is used with beam-based generation. length_penalty > 0.0 promotes longer sequences, while length_penalty < 0.0 encourages shorter sequences.
     bool past_present_share_buffer{};  // The past/present kv tensors are shared and allocated once to max_length (cuda only)
     int random_seed{-1};               // -1 = Seed with random device, otherwise use value to seed RNG
+    int chunk_size{-1};              // Chunk size for prefix chunking during context processing. -1 = disabled, >0 = enabled with specified chunk size.
   } search;
 
   void AddMapping(const std::string& nominal_name, const std::string& graph_name);

src/models/decoder_only.cpp

@@ -29,10 +29,10 @@ void DecoderOnly_State::SetExtraInputs(const std::vector<ExtraInput>& extra_inpu
 
 DeviceSpan<float> DecoderOnly_State::Run(int total_length, DeviceSpan<int32_t>& next_tokens, DeviceSpan<int32_t> next_indices) {
   size_t num_tokens = next_tokens.size();
-  const size_t chunk_size = 1024; // Experimental value
+  const size_t chunk_size = static_cast<size_t>(model_.config_->search.chunk_size);
 
-  if (num_tokens > chunk_size) {
-    // Chunking logic for context phase - process in chunks of 512 tokens
+  if (chunk_size > 0 && num_tokens > chunk_size) {
+    // Chunking logic for context phase - process in chunks based on configured chunk_size
     size_t processed_tokens = 0;
     int length = total_length - static_cast<int>(num_tokens);
     while (processed_tokens < num_tokens) {