fix: avoid suppressing EOS again after request recomputation

csrc/engine/prepare_sample.cpp

@@ -18,7 +18,8 @@ int num_prompt_logprobs(const SequenceMetadata &seq_metadata) {
 }
 
 bool suppresses_eos(const SequenceMetadata &seq_metadata) {
-    return seq_metadata.sampling.first_n_token_no_eos == 1 && seq_metadata.is_last_prompt_chunk();
+    return seq_metadata.sampling.first_n_token_no_eos == 1 && seq_metadata.is_last_prompt_chunk() &&
+           seq_metadata.seq_data->get_output_len() == 0;
 }
 
 void copy_to_default_device(torch::Tensor &tensor) { tensor = tensor.to(get_default_device(), /*non_blocking=*/true); }

tests/engine/test_eos_suppression.py

@@ -0,0 +1,123 @@
+import os
+from collections.abc import Generator
+
+import numpy as np
+import pytest
+import torch
+import torch.distributed as dist
+
+from deepinfer._C.engine import BlockManager, Scheduler, Sequence, SequenceMetadataPool, SequenceStatus, prepare_sampler_inputs
+from deepinfer.config import SchedulerConfig
+from deepinfer.engine.engine_type import EngineType
+from deepinfer.environ import rearm_all
+from deepinfer.sampling_params import SamplingParams
+
+pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]
+
+
+@pytest.fixture
+def process_group() -> Generator[dist.ProcessGroup]:
+    dist.init_process_group("gloo", store=dist.HashStore(), rank=0, world_size=1)
+    try:
+        yield dist.distributed_c10d._get_default_group()
+    finally:
+        dist.destroy_process_group()
+
+
+def make_sequence(seq_id: int, prompt: list[int], first_n_token_no_eos: int) -> Sequence:
+    tokens = np.asarray(prompt, dtype=np.int32)
+    sampling = SamplingParams(max_tokens=8, first_n_token_no_eos=first_n_token_no_eos)
+    sampling.eos_id = 1
+    return Sequence(
+        seq_id=seq_id,
+        prompt_token_ids=tokens,
+        compressed_prompt_token_ids=tokens,
+        multi_modal_data=None,
+        block_size=2,
+        request_id=str(seq_id),
+        user_id="test",
+        sampling_params=sampling,
+        arrival_time=0.0,
+        golden_token_ids=None,
+    )
+
+
+@pytest.mark.parametrize("first_n_token_no_eos", [0, 1])
+@pytest.mark.parametrize("chunked_prefill", [False, True])
+def test_eos_suppression_after_recompute(process_group: dist.ProcessGroup, first_n_token_no_eos: int, chunked_prefill: bool) -> None:
+    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
+    os.environ["DEEPINFER_ZERO_OVERHEAD_SCHEDULING"] = "0"
+    rearm_all()
+    config = SchedulerConfig(
+        max_num_batched_tokens=2 if chunked_prefill else 16,
+        max_num_seqs=1,
+        max_model_len=16,
+        chunked_prefill=chunked_prefill,
+        prefill_attempt_interval=1,
+        prefill_any_attempt_interval=1,
+    )
+    config.engine_type = EngineType.NORMAL.value
+    blocks = BlockManager(8)
+    scheduler = Scheduler(
+        block_size=2,
+        block_manager=blocks,
+        state_manager=None,
+        eos_token_id=1,
+        dp_group=process_group,
+        scheduler_config=config,
+        truncate_len=0,
+        lookahead_tokens=0,
+        prefill_max_sum_seqlen_k=16,
+        flops_estimator=None,
+        engram_lookup_table=[],
+        cp_size=1,
+        vocab_size=32,
+        log_stats=False,
+    )
+    pool = SequenceMetadataPool()
+    seq = make_sequence(0, [2, 3, 4], first_n_token_no_eos)
+    scheduler.add_sequence(seq)
+    empty = torch.empty(0, dtype=torch.int32)
+
+    # 完成首次 prefill，只有生成首个 token 的步骤允许屏蔽 EOS。
+    for _ in range(2 if chunked_prefill else 1):
+        metadata, outputs = scheduler.schedule(pool)
+        assert outputs.prompt_run
+        assert len(metadata) == 1
+        prepared = prepare_sampler_inputs(metadata, 0)
+        suppress = first_n_token_no_eos == 1 and metadata[0].is_last_prompt_chunk
+        assert prepared.driver.prefill.eos_suppression_rows == ([0] if suppress else [])
+        assert prepared.driver.prefill.eos_token_ids == ([1] if suppress else [])
+        scheduler.process_outputs(outputs, empty, torch.tensor([7], dtype=torch.int32), None)
+        pool.release(metadata)
+    assert seq.data.get_token_ids() == [2, 3, 4, 7]
+
+    # 通过实际分配缓存块耗尽可用容量，触发调度器抢占。
+    resident = make_sequence(1, [2] * 10, 0)
+    assert scheduler.try_allocate_seq(resident)
+    assert blocks.get_num_free_blocks() == 0
+    metadata, outputs = scheduler.schedule(pool)
+    assert outputs.is_empty()
+    assert metadata == []
+    assert scheduler.get_num_preempted_seqs() == 1
+    assert seq.status == SequenceStatus.WAITING
+    assert seq.num_cached_tokens == 0
+    assert seq.get_prompt_len() == 3
+    assert seq.data.get_token_ids() == [2, 3, 4, 7]
+    pool.release(metadata)
+    scheduler.free_seq(resident)
+
+    # 重算保留已生成的 token，最后一个分块也应允许生成 EOS。
+    for _ in range(2 if chunked_prefill else 1):
+        metadata, outputs = scheduler.schedule(pool)
+        assert outputs.prompt_run
+        assert len(metadata) == 1
+        prepared = prepare_sampler_inputs(metadata, 0)
+        assert prepared.driver.prefill.eos_suppression_rows == []
+        assert prepared.driver.prefill.eos_token_ids == []
+        scheduler.process_outputs(outputs, empty, torch.tensor([1], dtype=torch.int32), None)
+        pool.release(metadata)
+    assert seq.status == SequenceStatus.FINISHED_STOPPED
+    assert seq.data.get_token_ids() == [2, 3, 4, 7, 1]
+    scheduler.free_seq(seq)
+    assert blocks.get_num_free_blocks() == 7