diff --git a/csrc/engine/prepare_sample.cpp b/csrc/engine/prepare_sample.cpp index 836f50a..afa903e 100644 --- a/csrc/engine/prepare_sample.cpp +++ b/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); } diff --git a/tests/engine/test_eos_suppression.py b/tests/engine/test_eos_suppression.py new file mode 100644 index 0000000..3392cc3 --- /dev/null +++ b/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