mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
Merge 6a934611d538bcce54434c73a886459bca36b021 into 553741557a57df44a46fb3ab72b6aafeaf945337
This commit is contained in:
Md Jahidul Hosssain Mekat
GitHub
commit
fe10625dfc
@ -310,7 +310,11 @@ struct Scheduler {
|
||||
if (!is_any_point && !is_all_point) {
|
||||
return _schedule_decode(acquire);
|
||||
}
|
||||
bool want = (get_gpu_cache_usage() <= GPU_CACHE_PREFILL_WATERMARK) && _has_prefill_waiting();
|
||||
// A prompt prefilled in chunks reserves the blocks of its whole prompt with its first chunk, and only its remaining chunks can
|
||||
// give them back. Holding it to the watermark would stall it for as long as its own blocks keep the usage above it, which is
|
||||
// for good if it needs more than that share of the cache by itself.
|
||||
const bool continues_chunked_prefill = !waiting.empty() && waiting.front()->schedule.chunked_rest_len > 0;
|
||||
bool want = _has_prefill_waiting() && (continues_chunked_prefill || get_gpu_cache_usage() <= GPU_CACHE_PREFILL_WATERMARK);
|
||||
bool do_prefill;
|
||||
if (is_any_point) {
|
||||
do_prefill = (_dp_allreduce(want ? 0 : 1, c10d::ReduceOp::MIN) == 0); // ANY：任一 rank 想做即触发
|
||||
@ -362,6 +366,8 @@ struct Scheduler {
|
||||
// 入队前已判死的 seq 不跑 prefill，进 ignored_seqs 带 finish_reason 回传
|
||||
if (seq->schedule.status == SequenceStatus::FINISHED_ABORTED || seq->schedule.status == SequenceStatus::FINISHED_VE_FAILED ||
|
||||
seq->schedule.status == SequenceStatus::FINISHED_SO_FAILED) {
|
||||
// A chunked prefill that was aborted between chunks (or a fork) still holds KV blocks / a state slot.
|
||||
free_seq_nolock(*seq);
|
||||
ignored_seqs.push_back(seq);
|
||||
waiting.pop_front();
|
||||
continue;
|
||||
@ -569,6 +575,7 @@ struct Scheduler {
|
||||
std::vector<SequencePtr> ignored_seqs;
|
||||
for (auto it = waiting.begin(); it != waiting.end();) {
|
||||
if ((*it)->is_finished()) {
|
||||
free_seq_nolock(**it); // see _schedule_prefill: an aborted waiting seq may still hold KV blocks / a state slot
|
||||
ignored_seqs.push_back(std::move(*it));
|
||||
it = waiting.erase(it);
|
||||
} else {
|
||||
@ -735,10 +742,12 @@ struct Scheduler {
|
||||
return py::make_tuple(std::move(seq_metadata_list), std::move(p_scheduler_outputs));
|
||||
}
|
||||
|
||||
/// Release everything the sequence holds. Safe to call again: nothing is left to release the second time.
|
||||
void free_seq_nolock(Sequence &seq) {
|
||||
release_mtp_info_slot(seq);
|
||||
if (seq.cache.state_id != -1) {
|
||||
state_manager->free_blocks_nolock({seq.cache.state_id});
|
||||
seq.cache.state_id = -1;
|
||||
}
|
||||
block_manager->free_blocks_nolock(seq.cache.block_table);
|
||||
seq.cache.block_table.clear();
|
||||
@ -782,7 +791,15 @@ struct Scheduler {
|
||||
LOG_WARN("Request {} is preempted and will be recomputed.", seq->request.request_id);
|
||||
seq->reset_for_recompute(); // 复位 seq 自身的滚动状态（cache.num_tokens、inflight 计数、MTP 状态），清单见
|
||||
// Sequence::reset_for_recompute
|
||||
waiting.push_front(seq);
|
||||
// Queue it ahead of the fresh requests but behind the chunked prefills in progress. Those already hold the KV blocks of
|
||||
// their whole prompt and only need compute to finish. _schedule_prefill stops at the first sequence it cannot allocate, and
|
||||
// a preempted request usually does not fit right away (it needs a block more than it released): in front of them it would
|
||||
// keep them from ever being scheduled, and the blocks they hold from ever being released. A sequence that cannot be
|
||||
// scheduled yet (its multimodal inputs are still being processed) does not count as one of them, it would hold this
|
||||
// request up instead.
|
||||
waiting.insert(std::find_if(waiting.begin(), waiting.end(),
|
||||
[](const SequencePtr &s) { return s->cache.block_table.empty() || !s->schedule.schedulable; }),
|
||||
seq);
|
||||
seq->schedule.status = SequenceStatus::WAITING;
|
||||
free_seq_nolock(*seq);
|
||||
}
|
||||
|
||||
275
tests/engine/test_scheduler_resources.py
Normal file
275
tests/engine/test_scheduler_resources.py
Normal file
@ -0,0 +1,275 @@
|
||||
"""KV blocks and state slots must always find their way back to the managers, and the requests holding them must always be able to finish.
|
||||
|
||||
Drives the real C++ scheduler through the part of LLMEngine.step() that touches it: schedule, free the finished
|
||||
sequences, feed the (fake) model outputs back. No GPU or model is needed, only the compiled extension and a CPU
|
||||
process group.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Generator
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch # noqa: F401 Load torch libraries before the C++ extension.
|
||||
import torch.distributed as dist
|
||||
|
||||
from deepinfer._C.engine import BlockManager, Scheduler, Sequence, SequenceMetadataPool, process_MTPInfo
|
||||
from deepinfer.sampling_params import SamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]
|
||||
|
||||
EOS = 1
|
||||
VOCAB_SIZE = 1000
|
||||
GENERATED_TOKEN = 5 # never EOS, so requests end by max_tokens
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def cpu_group(tmp_path_factory: pytest.TempPathFactory) -> Generator[dist.ProcessGroup]:
|
||||
initialized_here = not dist.is_initialized()
|
||||
if initialized_here:
|
||||
rendezvous = Path(tmp_path_factory.mktemp("rendezvous")) / "store"
|
||||
dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=0, world_size=1)
|
||||
yield dist.new_group(ranks=[0], backend="gloo")
|
||||
if initialized_here:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
class _Engine:
|
||||
"""schedule -> free finished sequences -> process outputs, without speculative decoding or zero-overhead scheduling."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
group: dist.ProcessGroup,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
num_blocks: int,
|
||||
block_size: int,
|
||||
max_num_batched_tokens: int,
|
||||
num_states: int = 0,
|
||||
max_num_seqs: int = 2,
|
||||
prefill_attempt_interval: int = 2,
|
||||
) -> None:
|
||||
monkeypatch.setenv("DEEPINFER_USE_NEXTN_DECODE", "0") # read by the C++ scheduler when it is built
|
||||
self.num_blocks = num_blocks
|
||||
self.num_states = num_states
|
||||
self.block_size = block_size
|
||||
self.block_manager = BlockManager(num_blocks)
|
||||
self.state_manager = BlockManager(num_states) if num_states else None
|
||||
max_model_len = 64
|
||||
config = SimpleNamespace(
|
||||
max_model_len=max_model_len,
|
||||
runtime_max_model_len=max_model_len + 16,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
max_num_seqs=max_num_seqs,
|
||||
stop_preempted_request=False,
|
||||
chunked_prefill=True,
|
||||
engine_type=0, # EngineType.NORMAL
|
||||
prefill_attempt_interval=prefill_attempt_interval,
|
||||
prefill_any_attempt_interval=prefill_attempt_interval + 1,
|
||||
)
|
||||
self.scheduler = Scheduler(
|
||||
block_size, self.block_manager, self.state_manager, EOS, group, config, 0, 0, 2**30, None, [], 1, VOCAB_SIZE, False
|
||||
)
|
||||
self.pool = SequenceMetadataPool()
|
||||
self.sequences: dict[str, Sequence] = {}
|
||||
|
||||
def add_request(
|
||||
self, prompt_len: int, max_tokens: int, *, multi_modal_data: dict | None = None, allocate_blocks: bool = False
|
||||
) -> Sequence:
|
||||
request_id = str(len(self.sequences))
|
||||
prompt = np.full(prompt_len, 7, dtype=np.int32)
|
||||
seq = Sequence(
|
||||
seq_id=len(self.sequences), prompt_token_ids=prompt, compressed_prompt_token_ids=prompt, multi_modal_data=multi_modal_data,
|
||||
block_size=self.block_size, request_id=request_id, user_id="user",
|
||||
sampling_params=SamplingParams(max_tokens=max_tokens, eos_id=EOS), arrival_time=0.0, golden_token_ids=None,
|
||||
) # fmt: skip
|
||||
self.sequences[request_id] = seq
|
||||
if allocate_blocks:
|
||||
assert self.scheduler.try_allocate_seq(seq)
|
||||
self.scheduler.add_sequence(seq)
|
||||
return seq
|
||||
|
||||
def step(self, *, decode_only: bool = False) -> None:
|
||||
if decode_only:
|
||||
# In a data-parallel engine every rank has to agree to a prompt step, so a rank can run decode steps between two chunks
|
||||
# of a prompt it is prefilling. A decode_step that is no attempt point gives the same schedule on a single rank.
|
||||
self.scheduler.decode_step = 1
|
||||
seq_metadata_list, outputs = self.scheduler.schedule(self.pool)
|
||||
for seq in outputs.finished_seqs:
|
||||
self.scheduler.free_seq(seq)
|
||||
n = len(outputs.scheduled_seqs)
|
||||
process_MTPInfo(outputs, False, torch.zeros((n, 0), dtype=torch.int32), None)
|
||||
self.scheduler.process_outputs(
|
||||
outputs, torch.empty(0, dtype=torch.int32), torch.full((n,), GENERATED_TOKEN, dtype=torch.int32), None
|
||||
)
|
||||
self.pool.release(seq_metadata_list)
|
||||
|
||||
def run(self, max_steps: int = 300) -> None:
|
||||
for _ in range(max_steps):
|
||||
if not self.scheduler.has_unfinished_seqs():
|
||||
return
|
||||
self.step()
|
||||
raise AssertionError(f"no progress after {max_steps} steps: {self.scheduler.get_num_unfinished_seqs()} requests unfinished")
|
||||
|
||||
def assert_all_resources_returned(self) -> None:
|
||||
# Block 0 (and state slot 0) is reserved by the managers and never handed out.
|
||||
assert self.scheduler.get_num_free_gpu_blocks() == self.num_blocks - 1
|
||||
if self.state_manager is not None:
|
||||
assert self.state_manager.get_num_free_blocks() == self.num_states - 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_states", [0, 8], ids=["no_state", "with_state"])
|
||||
def test_abort_between_prefill_chunks_releases_the_sequence_resources(
|
||||
cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch, num_states: int
|
||||
) -> None:
|
||||
engine = _Engine(cpu_group, monkeypatch, num_blocks=17, block_size=4, max_num_batched_tokens=8, num_states=num_states)
|
||||
long_prompt = engine.add_request(prompt_len=30, max_tokens=3) # 4 chunks of at most 8 tokens
|
||||
engine.add_request(prompt_len=6, max_tokens=3)
|
||||
|
||||
engine.step()
|
||||
# The first chunk is done: the request waits for its next chunk and already owns the KV blocks of the whole prompt.
|
||||
assert long_prompt.status.name == "WAITING"
|
||||
assert long_prompt.num_cached_tokens == 8
|
||||
assert engine.scheduler.get_num_free_gpu_blocks() < engine.num_blocks - 1
|
||||
|
||||
engine.scheduler.abort_sequence([long_prompt.request_id])
|
||||
engine.run()
|
||||
|
||||
assert long_prompt.status.name == "FINISHED_ABORTED"
|
||||
engine.assert_all_resources_returned()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_states", [0, 8], ids=["no_state", "with_state"])
|
||||
def test_abort_during_a_decode_step_releases_the_sequence_resources(
|
||||
cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch, num_states: int
|
||||
) -> None:
|
||||
engine = _Engine(cpu_group, monkeypatch, num_blocks=17, block_size=4, max_num_batched_tokens=8, num_states=num_states)
|
||||
engine.add_request(prompt_len=6, max_tokens=12)
|
||||
engine.step()
|
||||
long_prompt = engine.add_request(prompt_len=30, max_tokens=3) # 4 chunks of at most 8 tokens
|
||||
engine.step()
|
||||
assert long_prompt.status.name == "WAITING"
|
||||
assert long_prompt.num_cached_tokens == 8
|
||||
|
||||
# The aborted prompt is dropped from the queue by a decode step this time, not by a prompt step.
|
||||
engine.scheduler.abort_sequence([long_prompt.request_id])
|
||||
engine.step(decode_only=True)
|
||||
assert long_prompt.status.name == "FINISHED_ABORTED"
|
||||
assert not any(seq.request_id == long_prompt.request_id for seq in engine.scheduler.waiting)
|
||||
|
||||
engine.run()
|
||||
engine.assert_all_resources_returned()
|
||||
|
||||
|
||||
def test_abort_before_the_first_chunk_leaves_nothing_to_release(cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
engine = _Engine(cpu_group, monkeypatch, num_blocks=17, block_size=4, max_num_batched_tokens=8, num_states=8)
|
||||
request = engine.add_request(prompt_len=30, max_tokens=3)
|
||||
engine.scheduler.abort_sequence([request.request_id])
|
||||
|
||||
engine.run()
|
||||
|
||||
assert request.status.name == "FINISHED_ABORTED"
|
||||
engine.assert_all_resources_returned()
|
||||
|
||||
|
||||
def _preempt_during_chunked_prefill(
|
||||
group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch, *, num_states: int = 0
|
||||
) -> tuple[_Engine, list[Sequence], Sequence]:
|
||||
"""Two requests decode while a long prompt is prefilled in chunks, and the cache is too small for all of them to keep growing.
|
||||
|
||||
The cache has 6 usable blocks of 4 tokens. The prompt (12 tokens, chunks of 4) reserves 3 of them with its first chunk, the
|
||||
decoding requests hold one each, and the last block goes to one of them for its next token. The other one needs a new block too,
|
||||
is preempted and has to be recomputed, which takes 2 blocks while only the 1 it just released is free.
|
||||
"""
|
||||
engine = _Engine(group, monkeypatch, num_blocks=7, block_size=4, max_num_batched_tokens=4, max_num_seqs=4, num_states=num_states)
|
||||
decoding = [engine.add_request(prompt_len=4, max_tokens=12), engine.add_request(prompt_len=4, max_tokens=12)]
|
||||
long_prompt = engine.add_request(prompt_len=12, max_tokens=2)
|
||||
for _ in range(10):
|
||||
if long_prompt.num_cached_tokens:
|
||||
break
|
||||
engine.step()
|
||||
assert long_prompt.status.name == "WAITING"
|
||||
assert long_prompt.num_cached_tokens == 4
|
||||
assert engine.scheduler.get_num_free_gpu_blocks() == 1
|
||||
|
||||
for _ in range(10):
|
||||
if engine.scheduler.get_num_preempted_seqs():
|
||||
break
|
||||
engine.step(decode_only=True)
|
||||
assert engine.scheduler.get_num_preempted_seqs() == 1
|
||||
assert sorted(seq.status.name for seq in decoding) == ["RUNNING", "WAITING"]
|
||||
assert long_prompt.status.name == "WAITING"
|
||||
return engine, decoding, long_prompt
|
||||
|
||||
|
||||
def test_abort_of_a_preempted_request_does_not_free_its_state_slot_twice(
|
||||
cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
engine, decoding, _ = _preempt_during_chunked_prefill(cpu_group, monkeypatch, num_states=8)
|
||||
preempted = next(seq for seq in decoding if seq.status.name == "WAITING")
|
||||
# Preemption already gave its state slot back; releasing the aborted sequence again must not release it a second time.
|
||||
engine.scheduler.abort_sequence([preempted.request_id])
|
||||
|
||||
engine.run()
|
||||
|
||||
assert preempted.status.name == "FINISHED_ABORTED"
|
||||
engine.assert_all_resources_returned()
|
||||
|
||||
|
||||
def test_preempted_request_does_not_block_a_chunked_prefill_that_holds_blocks(
|
||||
cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
engine, decoding, long_prompt = _preempt_during_chunked_prefill(cpu_group, monkeypatch)
|
||||
|
||||
# The preempted request cannot be recomputed before the long prompt has finished its prefill and released its blocks.
|
||||
engine.run()
|
||||
|
||||
assert [seq.status.name for seq in decoding] == ["FINISHED_LENGTH_CAPPED"] * 2
|
||||
assert long_prompt.status.name == "FINISHED_LENGTH_CAPPED"
|
||||
engine.assert_all_resources_returned()
|
||||
|
||||
|
||||
def test_preempted_request_is_not_queued_behind_a_sequence_that_cannot_be_scheduled(
|
||||
cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
# 7 usable blocks of 4 tokens. The two decoding requests need 3 blocks each in the end, the multimodal request holds 3.
|
||||
engine = _Engine(cpu_group, monkeypatch, num_blocks=8, block_size=4, max_num_batched_tokens=64, max_num_seqs=3)
|
||||
decoding = [engine.add_request(prompt_len=4, max_tokens=8), engine.add_request(prompt_len=4, max_tokens=8)]
|
||||
# It is not schedulable until its images are processed, but it owns its KV blocks (examples/benchmark.py allocates them up front).
|
||||
multimodal = engine.add_request(prompt_len=12, max_tokens=3, multi_modal_data={"image": 1}, allocate_blocks=True)
|
||||
assert not multimodal.schedulable
|
||||
|
||||
for _ in range(40):
|
||||
if all(seq.status.name == "FINISHED_LENGTH_CAPPED" for seq in decoding):
|
||||
break
|
||||
engine.step()
|
||||
|
||||
# The cache is too small for both, so one of them was preempted and recomputed. Queued behind the multimodal request, which
|
||||
# cannot be scheduled yet, it would wait for it forever.
|
||||
assert engine.scheduler.get_num_preempted_seqs() > 0
|
||||
assert [seq.status.name for seq in decoding] == ["FINISHED_LENGTH_CAPPED"] * 2
|
||||
|
||||
engine.scheduler.set_schedulable(multimodal.request_id)
|
||||
engine.run()
|
||||
assert multimodal.status.name == "FINISHED_LENGTH_CAPPED"
|
||||
engine.assert_all_resources_returned()
|
||||
|
||||
|
||||
def test_prompt_that_takes_nearly_the_whole_cache_is_prefilled_to_the_end(
|
||||
cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
# 10 usable blocks of 4 tokens. The first chunk reserves all of them, far above the usage at which new prompts are admitted.
|
||||
engine = _Engine(cpu_group, monkeypatch, num_blocks=11, block_size=4, max_num_batched_tokens=8)
|
||||
prompt = engine.add_request(prompt_len=37, max_tokens=2) # 5 chunks of at most 8 tokens
|
||||
|
||||
engine.step()
|
||||
assert prompt.status.name == "WAITING"
|
||||
assert engine.scheduler.get_num_free_gpu_blocks() == 0
|
||||
|
||||
engine.run()
|
||||
|
||||
assert prompt.status.name == "FINISHED_LENGTH_CAPPED"
|
||||
engine.assert_all_resources_returned()
|
||||
Loading…
x
Reference in New Issue
Block a user