DeepInfer/tests/engine/test_eos_suppression.py

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