[NPU][EAGLE3] Fix issue when Draft model generates EOS token (openvinotoolkit#3293)

<!-- Keep your pull requests (PRs) as atomic as possible. That increases
the likelihood that an individual PR won't be stuck because of adjacent
problems, merge conflicts, or code review.
Your merged PR is going to appear in the automatically generated release
notes on GitHub. So the clearer the title the better. -->
## Description
**Background**

Eagle3 uses speculative decoding where a draft model generates candidate
tokens that are validated by a target model. During this process, one
critical issue existed:

Sequence State Management: When the draft model generates an `EOS`
token, the `Sampler` immediately marks the sequence as `FINISHED`. This
causes `get_running_sequences()` to return empty, leading to assertion
failures in subsequent iterations if the target model rejects the `EOS`.

**Solution**
1. Restore Running State After Draft `EOS`
When draft model generates `EOS`, explicitly restore sequence to
`RUNNING` state after sampling.
2. Early Termination on Draft `EOS`
Break draft generation loop immediately upon detecting EOS:

<!-- Jira ticket number (e.g., 123). Delete if there's no ticket. -->
[CVS-180738](https://jira.devtools.intel.com/browse/CVS-180738)


## Checklist:
- [x] This PR follows GenAI Contributing guidelines. <!-- Always follow
https://github.com/openvinotoolkit/openvino.genai?tab=contributing-ov-file#contributing.
If there are deviations, explain what and why. -->
- [x] Tests have been updated or added to cover the new code. <!--
Specify exactly which tests were added or updated. If the change isn't
maintenance related, update the tests at
https://github.com/openvinotoolkit/openvino.genai/tree/master/tests or
explain in the description why the tests don't need an update. -->
- [x] This PR fully addresses the ticket. <!--- If not, explain clearly
what is covered and what is not. If follow-up pull requests are needed,
specify in the description. -->
- [x] I have made corresponding changes to the documentation. <!-- Run
github.com/\<username>/openvino.genai/actions/workflows/deploy_gh_pages.yml
on your fork with your branch as a parameter to deploy a test version
with the updated content. Replace this comment with the link to the
built docs. If the documentation is updated in a separate PR, clearly
specify it. -->

Author: GuoliangShiIntel

src/cpp/src/speculative_decoding/stateful/eagle3_strategy.cpp

@@ -195,13 +195,8 @@ std::vector<int64_t> Eagle3InferWrapperBase::sample_tokens(const ov::Tensor& log
     auto sequence_group = get_sequence_group();
     OPENVINO_ASSERT(sequence_group, "SequenceGroup not initialized");
 
-    OPENVINO_ASSERT(get_running_sequence_count() == 1,
-                    "Eagle3 currently only supports single sequence, got ",
-                    get_running_sequence_count(),
-                    " sequences");
-
     auto current_seq = get_current_sequence();
-    OPENVINO_ASSERT(current_seq, "No running sequence at index 0");
+    OPENVINO_ASSERT(current_seq, "No sequence at index 0");
 
     const size_t prev_generated_len = current_seq->get_generated_len();
     const size_t logits_seq_len = shape[1];
@@ -312,9 +307,9 @@ void Eagle3TargetWrapper::initialize_sequence(const ov::Tensor& input_ids, const
     TokenIds prompt_ids(ids_data, ids_data + seq_len);
     m_sequence_group = std::make_shared<SequenceGroup>(0, prompt_ids, config, 0);
 
-    OPENVINO_ASSERT(get_running_sequence_count() == 1,
+    OPENVINO_ASSERT(m_sequence_group->num_total_seqs() == 1,
                     "Expected single sequence after initialization, got ",
-                    get_running_sequence_count());
+                    m_sequence_group->num_total_seqs());
 }
 
 InferenceOutput Eagle3TargetWrapper::infer(const ov::Tensor& input_ids,
@@ -378,9 +373,9 @@ void Eagle3DraftWrapper::initialize_sequence(const ov::Tensor& input_ids, const
     TokenIds draft_prompt_ids(ids_data + 1, ids_data + total_len);
     m_sequence_group = std::make_shared<SequenceGroup>(1, draft_prompt_ids, config, 0);
 
-    OPENVINO_ASSERT(get_running_sequence_count() == 1,
+    OPENVINO_ASSERT(m_sequence_group->num_total_seqs() == 1,
                     "Expected single sequence after initialization, got ",
-                    get_running_sequence_count());
+                    m_sequence_group->num_total_seqs());
 }
 
 InferenceOutput Eagle3DraftWrapper::infer(const ov::Tensor& input_ids,
@@ -578,15 +573,18 @@ EncodedResults StatefulEagle3LLMPipeline::generate_tokens(const EncodedInputs& i
     while (!eos_reached && generated_tokens < config.max_new_tokens &&
            m_target->get_sequence_length() < m_prompt_length + config.max_new_tokens &&
            streaming_status == ov::genai::StreamingStatus::RUNNING) {
-        auto result = run_speculative_iteration(input_token_count, static_cast<int64_t>(config.eos_token_id));
+        auto result = run_speculative_iteration(input_token_count,
+                                                static_cast<int64_t>(config.eos_token_id),
+                                                generated_tokens,
+                                                config.max_new_tokens);
 
         streaming_status = stream_generated_tokens(streamer_ptr, result.validated_tokens);
 
         // Update statistics
         total_draft_generated += m_draft_iterations;
         total_draft_accepted += result.accepted_tokens_count;
+        generated_tokens += result.validated_tokens.size();
         eos_reached = result.eos_reached;
-        generated_tokens++;
 
         // Prepare for next iteration (hidden states are stored in sequence)
         input_token_count = result.next_window_size;
@@ -639,11 +637,13 @@ EncodedResults StatefulEagle3LLMPipeline::generate_tokens(const EncodedInputs& i
 
 StatefulEagle3LLMPipeline::SpeculativeResult StatefulEagle3LLMPipeline::run_speculative_iteration(
     size_t input_token_count,
-    int64_t eos_token_id) {
+    int64_t eos_token_id,
+    size_t current_generated_tokens,
+    size_t max_new_tokens) {
     SpeculativeResult result;
 
-    OPENVINO_ASSERT(m_target->get_running_sequence_count() == 1 && m_draft->get_running_sequence_count() == 1,
-                    "Eagle3 speculative iteration requires single sequence per model");
+    OPENVINO_ASSERT(m_target->get_sequence_group() && m_draft->get_sequence_group(),
+                    "Eagle3 speculative iteration requires initialized sequence groups");
 
     auto target_hidden_states = m_target->get_current_sequence()->get_hidden_state();
     OPENVINO_ASSERT(target_hidden_states && target_hidden_states.get_size() > 0,
@@ -670,8 +670,20 @@ StatefulEagle3LLMPipeline::SpeculativeResult StatefulEagle3LLMPipeline::run_spec
     // Append first token to target model (draft model already has it from sampler)
     m_target->append_tokens({first_draft_token});
 
+    // Check if first draft token is EOS - if so, no need to generate more draft tokens
+    bool draft_eos_reached = (first_draft_token == eos_token_id);
+
+    // IMPORTANT: If draft generated EOS, sampler will mark the sequence as FINISHED.
+    // However, we need to keep the draft sequence in RUNNING state because:
+    // 1. Target model may reject this EOS during validation
+    // 2. Next iteration needs draft sequence to be accessible via get_running_sequences()
+    // Only target model's EOS decision should truly end the generation.
+    if (draft_eos_reached) {
+        m_draft->get_current_sequence()->set_status(SequenceStatus::RUNNING);
+    }
+
     // Step 2: Generate additional draft tokens using internal hidden states
-    for (size_t i = 1; i < m_draft_iterations; ++i) {
+    for (size_t i = 1; i < m_draft_iterations && !draft_eos_reached; ++i) {
         InferContext more_ctx;
         more_ctx.input_token_count = 1;
         more_ctx.use_target_hidden = false;
@@ -685,16 +697,24 @@ StatefulEagle3LLMPipeline::SpeculativeResult StatefulEagle3LLMPipeline::run_spec
         // During validation, target model will retrieve tokens from its own sequence
         // so we need to speculatively add draft predictions here
         m_target->append_tokens({draft_token});
+
+        if (draft_token == eos_token_id) {
+            draft_eos_reached = true;
+            // Keep draft sequence in RUNNING state (same reason as above)
+            m_draft->get_current_sequence()->set_status(SequenceStatus::RUNNING);
+        }
     }
 
     // Step 3: Validate draft tokens with target model
 
-    const size_t validation_window_size = m_draft_iterations + 1;
+    // Validation window is based on actual draft tokens generated (may be less than m_draft_iterations if EOS hit)
+    const size_t actual_draft_tokens = draft_candidates.size();
+    const size_t validation_window_size = actual_draft_tokens + 1;
 
     InferContext val_ctx;
     val_ctx.input_token_count = validation_window_size;
     val_ctx.sample_count = validation_window_size;
-    val_ctx.num_tokens_to_validate = m_draft_iterations;
+    val_ctx.num_tokens_to_validate = actual_draft_tokens;
     auto val_result = m_target->forward(val_ctx);
 
     // Sampler validates draft tokens and returns accepted + new sampled token
@@ -703,8 +723,34 @@ StatefulEagle3LLMPipeline::SpeculativeResult StatefulEagle3LLMPipeline::run_spec
     // Result: [accepted_draft_tokens..., new_sampled_token]
     const size_t accepted_count = validated_tokens.size() - 1;
     const int64_t target_predicted_token = validated_tokens.back();
-    const size_t tokens_to_remove = m_draft_iterations - accepted_count;
-    const size_t total_accepted_tokens = validated_tokens.size();
+    size_t tokens_to_remove = actual_draft_tokens - accepted_count;
+    size_t total_accepted_tokens = validated_tokens.size();
+
+    // Check if accepting all validated tokens would exceed max_new_tokens
+    size_t tokens_after_accept = current_generated_tokens + validated_tokens.size();
+    if (tokens_after_accept > max_new_tokens) {
+        // Truncate to exactly max_new_tokens
+        size_t excess_tokens = tokens_after_accept - max_new_tokens;
+        OPENVINO_ASSERT(excess_tokens < validated_tokens.size(),
+                        "excess_tokens (",
+                        excess_tokens,
+                        ") must be less than validated_tokens.size() (",
+                        validated_tokens.size(),
+                        ")");
+        size_t tokens_to_keep = validated_tokens.size() - excess_tokens;
+
+        validated_tokens.resize(tokens_to_keep);
+        total_accepted_tokens = tokens_to_keep;
+
+        m_target->truncate_sequence(m_prompt_length + max_new_tokens);
+
+        // Adjust metrics to reflect actual tokens kept after truncation
+        auto& target_batch_sizes = m_target->get_raw_perf_metrics().m_batch_sizes;
+        OPENVINO_ASSERT(!target_batch_sizes.empty(), "batch_sizes should have been recorded by sampler");
+        target_batch_sizes.back() = tokens_to_keep;
+
+        tokens_to_remove = actual_draft_tokens - (tokens_to_keep - 1);  // -1 for the new target token
+    }
 
     // Step 4: Synchronize sequences and KV cache
     // Target model's sequence is already updated by Sampler
@@ -732,8 +778,8 @@ StatefulEagle3LLMPipeline::SpeculativeResult StatefulEagle3LLMPipeline::run_spec
     auto next_hidden = ov::Tensor(current_hidden, start_coord, end_coord);
     m_target->get_current_sequence()->update_hidden_state(next_hidden);
 
-    result.accepted_tokens_count = accepted_count;
-    result.next_window_size = accepted_count + 1;
+    result.accepted_tokens_count = total_accepted_tokens - 1;
+    result.next_window_size = total_accepted_tokens;
     result.validated_tokens = std::move(validated_tokens);
     result.eos_reached = (target_predicted_token == eos_token_id);
 

src/cpp/src/speculative_decoding/stateful/eagle3_strategy.hpp

@@ -86,26 +86,12 @@ class Eagle3InferWrapperBase {
         return m_sequence_group;
     }
 
-    void set_sequence_group(SequenceGroup::Ptr sequence_group) {
-        m_sequence_group = sequence_group;
-        if (m_sequence_group) {
-            OPENVINO_ASSERT(get_running_sequence_count() == 1,
-                            "Eagle3 only supports single sequence, got ",
-                            get_running_sequence_count());
-        }
-    }
-
-    /// @brief Returns number of running sequences in the group
-    size_t get_running_sequence_count() const {
-        return m_sequence_group ? m_sequence_group->get_running_sequences().size() : 0;
-    }
-
     /// @brief Returns sequence at given index with bounds checking
     /// @param index Sequence index (0 for top-1)
     /// @return Sequence pointer or nullptr if index out of bounds
     Sequence::Ptr get_sequence(size_t index) const {
         if (m_sequence_group) {
-            auto sequences = m_sequence_group->get_running_sequences();
+            const auto& sequences = m_sequence_group->get_sequences();
             if (index < sequences.size()) {
                 return sequences[index];
             }
@@ -236,7 +222,10 @@ class StatefulEagle3LLMPipeline : public StatefulSpeculativePipelineBase {
         std::vector<int64_t> validated_tokens;
     };
 
-    SpeculativeResult run_speculative_iteration(size_t token_count, int64_t eos_token_id);
+    SpeculativeResult run_speculative_iteration(size_t token_count,
+                                                int64_t eos_token_id,
+                                                size_t current_generated_tokens,
+                                                size_t max_new_tokens);
 
     std::unique_ptr<Eagle3DraftWrapper> m_draft;
     std::unique_ptr<Eagle3TargetWrapper> m_target;

tests/python_tests/test_stateful_speculative_decoding.py

@@ -211,6 +211,11 @@ def add(a, b):
 
 Question: Can you please add 2 and 3
 A:""",
+    ),
+    (
+        "Qwen/Qwen3-1.7B",
+        "AngelSlim/Qwen3-1.7B_eagle3",
+        "What is the capital of Ireland?/no_think",
     )
 ]