Adding non-CB pipeline for Speculative Decoding

samples/cpp/text_generation/speculative_decoding_lm.cpp

@@ -13,10 +13,15 @@ int main(int argc, char* argv[]) try {
 
     ov::genai::GenerationConfig config;
     config.max_new_tokens = 100;
-    // Speculative decoding generation parameters like `num_assistant_tokens` and `assistant_confidence_threshold` are mutually excluded
-    // add parameter to enable speculative decoding to generate `num_assistant_tokens` candidates by draft_model per iteration
-    config.num_assistant_tokens = 5;
-    // add parameter to enable speculative decoding to generate candidates by draft_model while candidate probability is higher than `assistant_confidence_threshold`
+    // Speculative decoding generation parameters like `num_assistant_tokens` and `assistant_confidence_threshold` are mutually excluded.
+    // Add parameter to enable speculative decoding to generate `num_assistant_tokens` candidates by draft_model per iteration.
+    // NOTE: ContinuousBatching backend uses `num_assistant_tokens` as is. Stateful backend uses `num_assistant_tokens`'s copy as initial
+    // value and adjusts it based on recent number of accepted tokens. If `num_assistant_tokens` is not set, it defaults to `5` for both
+    // backends.
+    config.num_assistant_tokens = 4;
+    // Add parameter to enable speculative decoding to generate candidates by draft_model while candidate probability is higher than
+    // `assistant_confidence_threshold`.
+    // NOTE: `assistant_confidence_threshold` is supported only by ContinuousBatching backend.
     // config.assistant_confidence_threshold = 0.4;
 
     std::string main_model_path = argv[1];
@@ -25,6 +30,8 @@ int main(int argc, char* argv[]) try {
 
     // User can run main and draft model on different devices.
     // Please, set device for main model in `LLMPipeline` constructor and in in `ov::genai::draft_model` for draft.
+    // CPU, GPU and NPU can be used. Please be aware that GPU is performant only with Continuous Batching pipeline, so it is not recommented
+    // to use it in conjuction with NPU or in configuration when main model doesn't work in Paged Attention mode.
     std::string main_device = "CPU", draft_device = "CPU";
 
     ov::genai::LLMPipeline pipe(

samples/python/text_generation/speculative_decoding_lm.py

 def test_perf_metrics(generation_config, prompt): 
 def test_perf_metrics(generation_config, prompt): 
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-# Copyright (C) 2024 Intel Corporation
+# Copyright (C) 2024-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 import argparse
@@ -19,8 +19,10 @@ def main():
     args = parser.parse_args()
 
     # User can run main and draft model on different devices.
-    # Please, set device for main model in `openvino_genai.LLMPipeline` constructor and in openvino_genai.draft_model` for draft.
-    main_device = 'CPU'  # GPU can be used as well
+    # Please, set device for main model in `openvino_genai.LLMPipeline` constructor and in `openvino_genai.draft_model` for draft.
+    # CPU, GPU and NPU can be used. Please be aware that GPU is performant only with Continuous Batching pipeline, so it is not
+    # recommented to use it in conjuction with NPU or in configuration when main model doesn't work in Paged Attention mode.
+    main_device = 'CPU'
     draft_device = 'CPU'
 
     draft_model = openvino_genai.draft_model(args.draft_model_dir, draft_device)
@@ -29,10 +31,15 @@ def main():
 
     config = openvino_genai.GenerationConfig()
     config.max_new_tokens = 100
-    # Speculative decoding generation parameters like `num_assistant_tokens` and `assistant_confidence_threshold` are mutually excluded
-    # add parameter to enable speculative decoding to generate `num_assistant_tokens` candidates by draft_model per iteration
-    config.num_assistant_tokens = 5
-    # add parameter to enable speculative decoding to generate candidates by draft_model while candidate probability is higher than `assistant_confidence_threshold`
+    # Speculative decoding generation parameters like `num_assistant_tokens` and `assistant_confidence_threshold` are mutually excluded.
+    # Add parameter to enable speculative decoding to generate `num_assistant_tokens` candidates by draft_model per iteration.
+    # NOTE: ContinuousBatching backend uses `num_assistant_tokens` as is. Stateful backend uses `num_assistant_tokens`'s copy as initial
+    # value and adjusts it based on recent number of accepted tokens. If `num_assistant_tokens` is not set, it defaults to `5` for both
+    # backends.
+    config.num_assistant_tokens = 4
+    # Add parameter to enable speculative decoding to generate candidates by draft_model while candidate probability is higher than
+    # `assistant_confidence_threshold`.
+    # NOTE: `assistant_confidence_threshold` is supported only by ContinuousBatching backend.
     # config.assistant_confidence_threshold = 0.4
 
     # Since the streamer is set, the results will be printed 

src/cpp/include/openvino/genai/generation_config.hpp

@@ -273,7 +273,10 @@ operator|(const StructuredOutputConfig::CompoundGrammar& lhs,
  *
  * Assisting generation parameters:
  * @param assistant_confidence_threshold the lower token probability of candidate to be validated by main model in case of dynamic strategy candidates number update.
+          NOTE: `assistant_confidence_threshold` is supported only by ContinuousBatching backend for Speculative Decode.
  * @param num_assistant_tokens the defined candidates number to be generated by draft model/prompt lookup in case of static strategy candidates number update.
+ *        NOTE: ContinuousBatching backend for Speculative Decode uses `num_assistant_tokens` as is. Stateful backend for Speculative Decode uses `num_assistant_tokens`'s
+ *        copy as initial value and adjusts it based on recent number of accepted tokens. If `num_assistant_tokens` is not set, it defaults to `5` for both backends.
  * @param max_ngram_size is maximum ngram to use when looking for matches in the prompt.
  *
  * @param structured_output_config if set, the output will be a string constrained by the specified json_schema, regex, or EBNF grammar.

src/cpp/src/continuous_batching/pipeline.cpp

@@ -19,16 +19,6 @@
 using namespace ov::genai;
 
 namespace {
-ov::genai::ModelDesc
-extract_draft_model_from_config(ov::AnyMap& config) {
-    ov::genai::ModelDesc draft_model;
-    if (config.find(utils::DRAFT_MODEL_ARG_NAME) != config.end()) {
-        draft_model = config.at(utils::DRAFT_MODEL_ARG_NAME).as<ov::genai::ModelDesc>();
-        config.erase(utils::DRAFT_MODEL_ARG_NAME);
-    }
-    return draft_model;
-}
-
 bool
 extract_prompt_lookup_from_config(ov::AnyMap& config) {
     bool res = false;
@@ -53,7 +43,7 @@ ContinuousBatchingPipeline::ContinuousBatchingPipeline( const std::filesystem::p
                                                         const ov::AnyMap& vision_encoder_properties) {
     auto start_time = std::chrono::steady_clock::now();
     auto properties_without_draft_model = properties;
-    auto draft_model_desr = extract_draft_model_from_config(properties_without_draft_model);
+    auto draft_model_desr = utils::extract_draft_model_from_config(properties_without_draft_model);
     auto is_prompt_lookup_enabled = extract_prompt_lookup_from_config(properties_without_draft_model);
 
     auto model = utils::read_model(models_path, properties);
@@ -92,7 +82,7 @@ ContinuousBatchingPipeline::ContinuousBatchingPipeline(
     const ov::AnyMap& properties) {
     auto start_time = std::chrono::steady_clock::now();
     auto properties_without_draft_model = properties;
-    auto draft_model_desr = extract_draft_model_from_config(properties_without_draft_model);
+    auto draft_model_desr = utils::extract_draft_model_from_config(properties_without_draft_model);
     auto is_prompt_lookup_enabled = extract_prompt_lookup_from_config(properties_without_draft_model);
 
     auto model = utils::read_model(models_path, properties_without_draft_model);
@@ -133,7 +123,7 @@ ContinuousBatchingPipeline::ContinuousBatchingPipeline(
     auto start_time = std::chrono::steady_clock::now();
 
     auto properties_without_draft_model = properties;
-    auto draft_model_desr = extract_draft_model_from_config(properties_without_draft_model);
+    auto draft_model_desr = utils::extract_draft_model_from_config(properties_without_draft_model);
     auto is_prompt_lookup_enabled = extract_prompt_lookup_from_config(properties_without_draft_model);
     auto model = utils::singleton_core().read_model(model_str, weights_tensor);
 
@@ -176,7 +166,7 @@ ContinuousBatchingPipeline::ContinuousBatchingPipeline(
     auto start_time = std::chrono::steady_clock::now();
 
     auto properties_without_draft_model = properties;
-    auto draft_model_desr = extract_draft_model_from_config(properties_without_draft_model);
+    auto draft_model_desr = utils::extract_draft_model_from_config(properties_without_draft_model);
     auto is_prompt_lookup_enabled = extract_prompt_lookup_from_config(properties_without_draft_model);
     auto model_pair = utils::get_model_weights_pair(models_map, "language");
     auto model = utils::singleton_core().read_model(model_pair.first, model_pair.second);

src/cpp/src/continuous_batching/timer.hpp

@@ -14,6 +14,8 @@ class ManualTimer {
 public:
     ManualTimer(const std::string& title) :
         m_total(0.),
+        m_start(),
+        m_end(),
         m_title(title) {
     }
 
@@ -42,6 +44,12 @@ class ManualTimer {
         return m_total;
     }
 
+    void clear() {
+        m_total = 0.0;
+        m_start = std::chrono::steady_clock::time_point();
+        m_end = std::chrono::steady_clock::time_point();
+    }
+
     ~ManualTimer() {
         // std::cout << m_title << ": " << m_total / 1e6 << " secs" << std::endl;
     }

src/cpp/src/llm/pipeline.cpp

@@ -9,10 +9,10 @@
 #include "openvino/genai/llm_pipeline.hpp"
 #include "openvino/genai/perf_metrics.hpp"
 
-#include "llm/pipeline_static.hpp"
 #include "llm/pipeline_stateful.hpp"
 #include "llm/pipeline_continuous_batching_adapter.hpp"
 #include "speculative_decoding/speculative_decoding_impl.hpp"
+#include "speculative_decoding/speculative_decoding_stateful.hpp"
 #include "utils.hpp"
 
 namespace ov {
@@ -60,6 +60,47 @@ std::pair<std::string, Any> draft_model(
     return { utils::DRAFT_MODEL_ARG_NAME, Any::make<ModelDesc>(model, tokenizer, device, plugin_config, scheduler_config, generation_config) };
 }
 
+class StatefulPipeline {
+public:
+static std::unique_ptr<LLMPipelineImplBase> create(
+    const std::filesystem::path& models_path,
+    const ov::genai::Tokenizer& tokenizer,
+    const std::string& device,
+    const ov::AnyMap& properties) {
+    return create(
+        ov::genai::utils::read_model(models_path, properties),
+        tokenizer,
+        device,
+        properties,
+        utils::from_config_json_if_exists(models_path));
+}
+
+static std::unique_ptr<LLMPipelineImplBase> create(
+    const std::filesystem::path& models_path,
+    const std::string& device,
+    const ov::AnyMap& plugin_config) {
+    return create(models_path, Tokenizer(models_path, plugin_config), device, plugin_config);
+}
+
+static std::unique_ptr<LLMPipelineImplBase> create(
+    const std::shared_ptr<ov::Model>& model,
+    const ov::genai::Tokenizer& tokenizer,
+    const std::string& device,
+    const ov::AnyMap& properties,
+    const ov::genai::GenerationConfig& generation_config) {
+
+    auto properties_without_draft_model = properties;
+    auto draft_model_descr = ov::genai::utils::extract_draft_model_from_config(properties_without_draft_model);
+    if (draft_model_descr.model != nullptr) {
+        auto main_model_descr = ov::genai::ModelDesc(model, tokenizer, device, properties_without_draft_model, {}, generation_config);
+        return std::make_unique<StatefulSpeculativeLLMPipeline>(main_model_descr, draft_model_descr);
+    }
+
+    return std::make_unique<StatefulLLMPipeline>(model, tokenizer, device,
+        properties_without_draft_model, generation_config);
+}
+};
+
 // Public LLMPipeline
 
 ov::genai::LLMPipeline::LLMPipeline(
@@ -80,14 +121,12 @@ ov::genai::LLMPipeline::LLMPipeline(
     auto start_time = std::chrono::steady_clock::now();
     auto [properties, attention_backend] = utils::extract_attention_backend(user_properties);
 
-    // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
-    if (utils::explicitly_requires_paged_attention(user_properties)) {
+    if (ov::genai::utils::is_npu_requested(device, properties)) {
+        m_pimpl = StatefulPipeline::create(models_path, tokenizer, device, properties);
+    } else if (utils::explicitly_requires_paged_attention(user_properties)) {
+        // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
         auto [device_properties, scheduler_config] = utils::extract_scheduler_config(properties, utils::get_latency_oriented_scheduler_config());
         m_pimpl = std::make_unique<ContinuousBatchingAdapter>(models_path, tokenizer, scheduler_config, device, device_properties);
-    } else if (device == "NPU") {
-        m_pimpl = properties.count("STATIC_PIPELINE")
-            ? static_llm::LLMPipelineFactory::create(models_path, tokenizer, properties)
-            : std::make_unique<StatefulLLMPipeline>(models_path, tokenizer, device, properties);
     } else if (attention_backend == PA_BACKEND) {
         // try to call CB adapter one more time, but with safe guard to silent exception
         try {
@@ -102,7 +141,7 @@ ov::genai::LLMPipeline::LLMPipeline(
     }
 
     if (m_pimpl == nullptr) {
-        m_pimpl = std::make_unique<StatefulLLMPipeline>(models_path, tokenizer, device, properties);
+        m_pimpl = StatefulPipeline::create(models_path, tokenizer, device, properties);
     }
 
     m_pimpl->save_load_time(start_time);
@@ -117,14 +156,12 @@ ov::genai::LLMPipeline::LLMPipeline(
 
     auto [properties, attention_backend] = utils::extract_attention_backend(user_properties);
 
-    // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
-    if (utils::explicitly_requires_paged_attention(user_properties)) {
+    if (ov::genai::utils::is_npu_requested(device, properties)) {
+        m_pimpl = StatefulPipeline::create(models_path, device, properties);
+    } else if (utils::explicitly_requires_paged_attention(user_properties)) {
+        // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
         auto [device_properties, scheduler_config] = utils::extract_scheduler_config(properties, utils::get_latency_oriented_scheduler_config());
         m_pimpl = std::make_unique<ContinuousBatchingAdapter>(models_path, scheduler_config, device, device_properties);
-    } else if (device == "NPU") {
-        m_pimpl = properties.count("STATIC_PIPELINE")
-            ? static_llm::LLMPipelineFactory::create(models_path, properties)
-            : std::make_unique<StatefulLLMPipeline>(models_path, device, properties);
     } else if (attention_backend == PA_BACKEND) {
         // try to call CB adapter one more time, but with safe guard to silent exception
         try {
@@ -139,7 +176,7 @@ ov::genai::LLMPipeline::LLMPipeline(
     }
 
     if (m_pimpl == nullptr) {
-        m_pimpl = std::make_unique<StatefulLLMPipeline>(models_path, device, properties);
+        m_pimpl = StatefulPipeline::create(models_path, device, properties);
     }
 
     m_pimpl->save_load_time(start_time);
@@ -157,24 +194,18 @@ ov::genai::LLMPipeline::LLMPipeline(
 
     auto [properties, attention_backend] = utils::extract_attention_backend(user_properties);
 
-    // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
-    if (utils::explicitly_requires_paged_attention(user_properties)) {
+    if (ov::genai::utils::is_npu_requested(device, properties)) {
+        m_pimpl = StatefulPipeline::create(
+            utils::singleton_core().read_model(model_str, weights_tensor),
+            tokenizer,
+            device,
+            properties,
+            generation_config);
+    } else if (utils::explicitly_requires_paged_attention(user_properties)) {
+        // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
         auto [device_properties, scheduler_config] = utils::extract_scheduler_config(properties, utils::get_latency_oriented_scheduler_config());
         m_pimpl = std::make_unique<ContinuousBatchingAdapter>(model_str, weights_tensor,
                                                               tokenizer, scheduler_config, device, device_properties, generation_config);
-    } else if (device == "NPU") {
-        m_pimpl = properties.count("STATIC_PIPELINE")
-            ? static_llm::LLMPipelineFactory::create(
-                  utils::singleton_core().read_model(model_str, weights_tensor),
-                  tokenizer,
-                  properties,
-                  generation_config)
-            : std::make_unique<StatefulLLMPipeline>(
-                utils::singleton_core().read_model(model_str, weights_tensor),
-                tokenizer,
-                device,
-                properties,
-                generation_config);
     } else if (attention_backend == PA_BACKEND) {
         // try to call CB adapter one more time, but with safe guard to silent exception
         try {
@@ -190,7 +221,7 @@ ov::genai::LLMPipeline::LLMPipeline(
     }
 
     if (m_pimpl == nullptr) {
-        m_pimpl = std::make_unique<StatefulLLMPipeline>(
+        m_pimpl = StatefulPipeline::create(
             utils::singleton_core().read_model(model_str, weights_tensor),
             tokenizer,
             device,

src/cpp/src/speculative_decoding/continuous_batching_for_speculative_decoding_impl.cpp

@@ -14,6 +14,11 @@ ContinuousBatchingPipeline::ContinuousBatchingForSpeculativeDecodingImpl::Contin
     bool is_validation_mode_enabled) {
     m_tokenizer = tokenizer;
     m_generation_config = generation_config;
+    if (m_generation_config.assistant_confidence_threshold == 0.f) {
+        if (m_generation_config.num_assistant_tokens == 0) {
+            m_generation_config.num_assistant_tokens = default_num_assistant_tokens;
+        }
+    }
     m_is_validation_mode_enabled = is_validation_mode_enabled;
     initialize_pipeline(model, scheduler_config, device, plugin_config);
 }
@@ -319,7 +324,7 @@ void ContinuousBatchingPipeline::ContinuousBatchingForSpeculativeDecodingImpl::m
         auto pipeline_metrics = get_metrics();
         if (num_generated_tokens > 0) {
             raw_perf_metrics.m_durations.emplace_back(generation_duration);
-            raw_perf_metrics.m_inference_durations[0] = MicroSeconds(pipeline_metrics.inference_duration);
+            raw_perf_metrics.m_inference_durations[0] += MicroSeconds(pipeline_metrics.inference_duration);
             raw_perf_metrics.m_batch_sizes.emplace_back(num_generated_tokens);
         }
 

src/cpp/src/speculative_decoding/continuous_batching_for_speculative_decoding_impl.hpp

@@ -11,6 +11,8 @@
 namespace ov::genai {
 class ContinuousBatchingPipeline::ContinuousBatchingForSpeculativeDecodingImpl : public ContinuousBatchingPipeline::ContinuousBatchingImpl {
 public:
+    const std::size_t default_num_assistant_tokens = 5;
+
     ContinuousBatchingForSpeculativeDecodingImpl() = default;
 
     ContinuousBatchingForSpeculativeDecodingImpl(const std::shared_ptr<ov::Model>& model,