mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-30 07:47:10 +08:00
147 lines
4.7 KiB
Python
147 lines
4.7 KiB
Python
"""CPU fixtures for the native scheduler; no model weights or device execution."""
|
|
|
|
from collections.abc import Callable, Iterator
|
|
from itertools import count
|
|
from types import SimpleNamespace
|
|
from typing import TYPE_CHECKING
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
if TYPE_CHECKING:
|
|
from torch.distributed import ProcessGroup
|
|
|
|
from deepinfer._C.engine import BlockManager, Scheduler, Sequence, SequenceMetadataPool
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def scheduler_process_group() -> Iterator["ProcessGroup"]:
|
|
import torch.distributed as dist
|
|
|
|
owns_group = not dist.is_initialized()
|
|
if owns_group:
|
|
dist.init_process_group("gloo", store=dist.HashStore(), rank=0, world_size=1)
|
|
elif dist.get_world_size() != 1 or dist.get_backend() != "gloo":
|
|
pytest.skip("Native scheduler unit tests require a single-rank CPU Gloo group")
|
|
try:
|
|
yield dist.group.WORLD
|
|
finally:
|
|
if owns_group:
|
|
dist.destroy_process_group()
|
|
|
|
|
|
@pytest.fixture
|
|
def scheduler_factory(
|
|
scheduler_process_group: "ProcessGroup", monkeypatch: pytest.MonkeyPatch
|
|
) -> Callable[..., tuple["Scheduler", "BlockManager", "BlockManager | None"]]:
|
|
from deepinfer._C.engine import BlockManager, Scheduler
|
|
|
|
def make(
|
|
num_blocks: int = 33,
|
|
*,
|
|
num_states: int | None = None,
|
|
stop_preempted_request: bool = True,
|
|
nextn: int = 0,
|
|
max_num_batched_tokens: int = 64,
|
|
max_num_seqs: int = 4,
|
|
max_model_len: int = 256,
|
|
) -> tuple["Scheduler", "BlockManager", "BlockManager | None"]:
|
|
monkeypatch.setenv("DEEPINFER_USE_NEXTN_DECODE", str(nextn))
|
|
blocks = BlockManager(num_blocks)
|
|
states = BlockManager(num_states) if num_states is not None else None
|
|
config = SimpleNamespace(
|
|
max_model_len=max_model_len,
|
|
runtime_max_model_len=max_model_len + 64,
|
|
max_num_batched_tokens=max_num_batched_tokens,
|
|
max_num_seqs=max_num_seqs,
|
|
stop_preempted_request=stop_preempted_request,
|
|
chunked_prefill=True,
|
|
engine_type=0,
|
|
prefill_attempt_interval=2,
|
|
prefill_any_attempt_interval=100,
|
|
)
|
|
scheduler = Scheduler(
|
|
block_size=16,
|
|
block_manager=blocks,
|
|
state_manager=states,
|
|
eos_token_id=2,
|
|
dp_group=scheduler_process_group,
|
|
scheduler_config=config,
|
|
truncate_len=0,
|
|
lookahead_tokens=0,
|
|
prefill_max_sum_seqlen_k=4096,
|
|
flops_estimator=None,
|
|
engram_lookup_table=[],
|
|
cp_size=1,
|
|
vocab_size=100,
|
|
log_stats=False,
|
|
)
|
|
return scheduler, blocks, states
|
|
|
|
return make
|
|
|
|
|
|
@pytest.fixture
|
|
def sequence_factory() -> Callable[..., "Sequence"]:
|
|
from deepinfer._C.engine import Sequence
|
|
from deepinfer.sampling_params import SamplingParams
|
|
|
|
ids = count()
|
|
|
|
def make(length: int = 128, *, request_id: str = "request", block_size: int = 16) -> "Sequence":
|
|
tokens = np.ones(length, dtype=np.int32)
|
|
return Sequence(
|
|
seq_id=next(ids),
|
|
prompt_token_ids=tokens,
|
|
compressed_prompt_token_ids=tokens,
|
|
multi_modal_data=None,
|
|
block_size=block_size,
|
|
request_id=request_id,
|
|
user_id="test",
|
|
sampling_params=SamplingParams(max_tokens=4, ignore_eos=True, eos_id=2, seq_seed=0),
|
|
arrival_time=0.0,
|
|
golden_token_ids=None,
|
|
)
|
|
|
|
return make
|
|
|
|
|
|
@pytest.fixture
|
|
def metadata_pool() -> "SequenceMetadataPool":
|
|
from deepinfer._C.engine import SequenceMetadataPool
|
|
|
|
return SequenceMetadataPool()
|
|
|
|
|
|
@pytest.fixture
|
|
def unscheduled_outputs() -> Callable:
|
|
import torch
|
|
|
|
from deepinfer._C.engine import RequestOutput, RequestOutputPool, build_request_outputs
|
|
|
|
def build(outputs: object) -> list[object]:
|
|
pool = RequestOutputPool()
|
|
pointers = build_request_outputs(
|
|
outputs,
|
|
64,
|
|
torch.empty(0, dtype=torch.int32),
|
|
torch.empty(0, dtype=torch.float32),
|
|
top_logprobs_=None,
|
|
top_token_ids_=None,
|
|
sum_smaller_probs_=None,
|
|
accepted_token_ids_=None,
|
|
accepted_token_logprobs_=None,
|
|
accepted_top_logprobs_=None,
|
|
accepted_top_token_ids_=None,
|
|
accepted_sum_smaller_probs_=None,
|
|
accepted_token_last_logprobs_=None,
|
|
pushed_tokens_before_finish=np.empty(0, dtype=np.int32),
|
|
prompt_logprobs=None,
|
|
request_output_pool=pool,
|
|
)
|
|
result = [RequestOutput.clone(int(pointer)) for pointer in pointers]
|
|
pool.release(pointers)
|
|
return result
|
|
|
|
return build
|