diff --git a/benchmarks/preemption.py b/benchmarks/preemption.py new file mode 100644 index 00000000..e1f53575 --- /dev/null +++ b/benchmarks/preemption.py @@ -0,0 +1,411 @@ +"""Paired local-preemption traffic benchmarks; profiling runs are separate from timing.""" + +import argparse +import hashlib +import itertools +import json +import math +import random +import statistics +import time +import traceback +from pathlib import Path + +CASES = [ + "_".join(parts) + for parts in itertools.product(("equal", "mixed"), ("uniform", "mixed"), ("low", "high")) +] + + +def workload(case, *, count, slots, seed, rate, vocab_size): + priority, lengths, _ = case.split("_") + rng = random.Random(seed) + requests = [] + arrival = 0.0 + for index in range(count): + short = lengths == "mixed" and index >= slots and index % 2 == 0 + prompt, output = (32, 16) if short else ((256, 128) if lengths == "mixed" else (128, 64)) + requests.append( + dict( + request_id=f"r{index:04d}", + arrival_s=arrival, + prompt_token_ids=[rng.randrange(2, vocab_size) for _ in range(prompt)], + max_tokens=output, + priority=0 if priority == "equal" or (index >= slots and index % 2 == 0) else 10, + ) + ) + arrival += rng.uniform(0.75, 1.25) / rate + return requests + + +def percentile(values, fraction): + values = sorted(values) + position = (len(values) - 1) * fraction + lo, hi = math.floor(position), math.ceil(position) + return values[lo] + (values[hi] - values[lo]) * (position - lo) + + +def metrics(records): + def summarize(rows): + return { + name: dict( + mean=statistics.mean(values), + p50=statistics.median(values), + p95=percentile(values, 0.95), + ) + for name in ("ttft_s", "e2e_s", "enqueue_ttft_s", "enqueue_e2e_s") + if (values := [row[name] for row in rows]) + } + + return dict( + all=summarize(records), + by_priority={ + str(priority): summarize([row for row in records if row["priority"] == priority]) + for priority in sorted({row["priority"] for row in records}) + }, + by_length={ + str(length): summarize( + [row for row in records if len(row["prompt_token_ids"]) == length] + ) + for length in sorted({len(row["prompt_token_ids"]) for row in records}) + }, + ) + + +def drive(engine, requests, state): + from vllm_rlt import SamplingParams + + start = time.perf_counter() + pending = iter(requests) + next_request = next(pending, None) + records = {} + state.update(requests=[], max_waiting=0, steps=0) + while next_request is not None or engine.has_unfinished_requests(): + now = time.perf_counter() - start + while next_request is not None and next_request["arrival_s"] <= now: + row = dict(next_request, enqueued_s=time.perf_counter() - start) + records[row["request_id"]] = row + state["requests"].append(row) + engine.add_request( + row["request_id"], + row["prompt_token_ids"], + SamplingParams( + max_tokens=row["max_tokens"], + priority=row["priority"], + ignore_eos=True, + ), + ) + next_request = next(pending, None) + state["max_waiting"] = max( + state["max_waiting"], + sum(request.stage.value == "waiting" for request in engine.scheduler.requests.values()), + ) + if engine.has_unfinished_requests(): + for output in engine.step(): + row = records[output.request_id] + delivered = time.perf_counter() - start + if output.token_ids and "first_output_s" not in row: + row["first_output_s"] = delivered + if output.finished: + row.update( + finished_s=delivered, + token_ids=output.token_ids, + exit_depths=output.exit_depths, + finish_reason=output.finish_reason, + ) + row["ttft_s"] = row["first_output_s"] - row["arrival_s"] + row["e2e_s"] = delivered - row["arrival_s"] + row["enqueue_ttft_s"] = row["first_output_s"] - row["enqueued_s"] + row["enqueue_e2e_s"] = delivered - row["enqueued_s"] + state["steps"] += 1 + elif next_request is not None: + time.sleep(min(0.002, max(0, next_request["arrival_s"] - now))) + engine.model_runner.synchronize() + state.update( + wall_s=time.perf_counter() - start, + preemptions=engine.preemption.preemptions, + resumptions=engine.preemption.resumptions, + kv_used_blocks=engine.cache_manager.num_used_blocks, + suspended_requests=len(engine.preemption.snapshots), + ) + if state["kv_used_blocks"] or state["suspended_requests"]: + raise AssertionError("completed traffic retained request KV or suspended state") + for row in state["requests"]: + if len(row["token_ids"]) != row["max_tokens"] or row["finish_reason"] != "length": + raise AssertionError("traffic did not produce the fixed output length") + state["tokens_per_s"] = ( + sum(len(row["token_ids"]) for row in state["requests"]) / state["wall_s"] + ) + state["metrics"] = metrics(state["requests"]) + + +def checkpoint_identity(folder): + names = [folder / "config.json", *sorted(folder.glob("*.safetensors"))] + if (index := folder / "model.safetensors.index.json").is_file(): + names.append(index) + if len(names) < 2: + raise ValueError("a complete local safetensors checkpoint is required") + result = {} + for path in names: + before = path.stat() + digest = hashlib.sha256() + with path.open("rb") as stream: + while block := stream.read(1024 * 1024): + digest.update(block) + after = path.stat() + fields = ("st_dev", "st_ino", "st_size", "st_mtime_ns", "st_ctime_ns") + if any(getattr(before, field) != getattr(after, field) for field in fields): + raise RuntimeError(f"checkpoint changed while hashing: {path.name}") + result[path.name] = dict(size=after.st_size, sha256=digest.hexdigest()) + return result + + +def trial(model, args, requests, output_dir, label, state, raw): + import torch + + from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + ProfileConfig, + SchedulerConfig, + SpeculativeConfig, + ) + from vllm_rlt.engine.llm_engine import LLMEngine + + engine = None + try: + engine = LLMEngine( + model, + attention_backend=args.attention_backend, + cache_config=CacheConfig( + kv_cache_memory_bytes=args.kv_gib * 1024**3, incremental_allocation=True + ), + scheduler_config=SchedulerConfig( + max_num_seqs=args.max_num_seqs, + max_num_batched_tokens=64, + prefill_chunk_size=64, + policy="priority", + enable_preemption=state["enable_preemption"], + ), + execution_config=ExecutionConfig(cuda_graphs=args.cuda_graphs), + speculative_config=SpeculativeConfig(3) if args.decode == "speculative" else None, + ) + # Same token workloads warm every arm before measurement; arrival time is + # zero only for warmup, so low-load idle time is not multiplied by warmup. + for index in range(args.warmups): + warm = dict(label=f"{label}/warmup-{index}", kind="warmup") + drive(engine, [dict(row, arrival_s=0.0) for row in requests], warm) + raw.write(json.dumps(warm) + "\n") + raw.flush() + engine.preemption.preemptions = engine.preemption.resumptions = 0 + if args.device == "cuda": + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + if state["kind"] == "profile": + engine.start_profile( + ProfileConfig( + output_dir=str(output_dir / "profiles" / label), + record_shapes=True, + with_stack=True, + profile_memory=True, + warmup=1, + active=args.profile_steps, + ), + scheduled=True, + ) + drive(engine, requests, state) + finally: + if engine is not None: + engine.close() # Stop capture and naturally join the artifact worker. + if state["kind"] == "profile": + state["profile"] = engine.profile_status() + if args.device == "cuda": + state["peak_allocated_bytes"] = torch.cuda.max_memory_allocated() + state["peak_reserved_bytes"] = torch.cuda.max_memory_reserved() + if state["kind"] == "profile" and not state["profile"]["success"]: + raise RuntimeError("profiling artifacts failed; inspect the retained profile status") + + +def paired_summary(rows): + result = {} + for case in sorted({row["case"] for row in rows}): + trials = [row for row in rows if row["case"] == case and row["kind"] == "measurement"] + if not trials: + continue + pairs = {} + for row in trials: + pairs.setdefault(row["pair"], {})[row["enable_preemption"]] = row + ratios = {} + for scope in trials[0]["metrics"]: + groups = [None] if scope == "all" else list(trials[0]["metrics"][scope]) + for group in groups: + for metric, statistic in itertools.product( + ("ttft_s", "e2e_s"), ("mean", "p50", "p95") + ): + + def value(row): + data = row["metrics"][scope] + return (data if group is None else data[group])[metric][statistic] + + values = [value(pair[True]) / value(pair[False]) for pair in pairs.values()] + ratios[f"{scope}/{group}/{metric}/{statistic}"] = dict( + per_pair=values, + mean=statistics.mean(values), + median=statistics.median(values), + ) + result[case] = dict(pairs=len(pairs), on_over_off=ratios) + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True, help="existing local Ouro checkpoint") + parser.add_argument("--output", type=Path, required=True, help="new output directory") + parser.add_argument("--low-rps", type=float, required=True) + parser.add_argument("--high-rps", type=float, required=True) + parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda") + parser.add_argument("--attention-backend", default="triton") + parser.add_argument("--decode", choices=("native", "speculative"), default="speculative") + parser.add_argument("--cuda-graphs", action="store_true") + parser.add_argument("--max-num-seqs", type=int, default=4) + parser.add_argument("--requests", type=int, default=32) + parser.add_argument("--pairs", type=int, default=5) + parser.add_argument("--warmups", type=int, default=1) + parser.add_argument("--kv-gib", type=int, default=2) + parser.add_argument("--seed", type=int, default=1729) + parser.add_argument("--case", choices=CASES, action="append") + parser.add_argument("--profile-case", choices=CASES) + parser.add_argument("--profile-only", action="store_true", help="skip the timing matrix") + parser.add_argument("--profile-steps", type=int, default=128) + parser.add_argument( + "--plan-only", action="store_true", help="print matrix without importing Torch" + ) + args = parser.parse_args() + if not ( + math.isfinite(args.low_rps) + and math.isfinite(args.high_rps) + and 0 < args.low_rps < args.high_rps + ): + parser.error("finite positive low-rps < high-rps required") + if ( + min( + args.max_num_seqs, + args.requests, + args.pairs, + args.warmups, + args.kv_gib, + args.profile_steps, + ) + < 1 + ): + parser.error("counts and capacities must be positive") + if args.requests <= args.max_num_seqs: + parser.error("requests must exceed max-num-seqs to exercise dynamic arrivals") + if args.profile_only and args.profile_case is None: + parser.error("profile-only requires profile-case") + cases = list(dict.fromkeys(args.case or CASES)) + plan = dict( + cases=cases, + low_rps=args.low_rps, + high_rps=args.high_rps, + pairs=args.pairs, + requests=args.requests, + seed=args.seed, + decode=args.decode, + timing="scheduled arrival to CPU output delivery; warmup/loading excluded", + default_enable_preemption=False, + scope="controlled local traffic; no quality score", + arguments={ + key: str(value) if isinstance(value, Path) else value + for key, value in vars(args).items() + }, + ) + if args.plan_only: + print(json.dumps(plan, indent=2)) + return + if not args.model.is_dir(): + parser.error("model must be an existing local checkpoint directory") + args.output.mkdir(parents=True, exist_ok=False) + identity = checkpoint_identity(args.model) + import torch + + from vllm_rlt.models import OuroForCausalLM + + torch.set_num_threads(1) + torch.backends.cuda.matmul.allow_tf32 = False + model = OuroForCausalLM.from_pretrained( + args.model, + device=args.device, + dtype=torch.bfloat16 if args.device == "cuda" else torch.float32, + ) + plan.update( + checkpoint=identity, + torch=torch.__version__, + source_sha256={ + str(path.relative_to(Path(__file__).resolve().parents[1])): hashlib.sha256( + path.read_bytes() + ).hexdigest() + for path in [ + Path(__file__).resolve(), + *Path(__file__).resolve().parents[1].joinpath("vllm_rlt").rglob("*.py"), + ] + }, + device_name=torch.cuda.get_device_name() if args.device == "cuda" else "CPU", + model_config=model.config.to_dict(), + ) + (args.output / "plan.json").write_text(json.dumps(plan, indent=2) + "\n") + rows = [] + with (args.output / "raw.jsonl").open("x") as raw: + runs = [] if args.profile_only else [(case, "measurement") for case in cases] + if args.profile_case: + runs.append((args.profile_case, "profile")) + for case, kind in runs: + for pair in range(1 if kind == "profile" else args.pairs): + rate = args.low_rps if case.endswith("_low") else args.high_rps + requests = workload( + case, + count=args.requests, + slots=args.max_num_seqs, + seed=args.seed + pair, + rate=rate, + vocab_size=model.config.vocab_size, + ) + outputs = {} + for enabled in (False, True) if pair % 2 == 0 else (True, False): + label = f"{kind}-{case}-{pair}-{'on' if enabled else 'off'}" + state = dict( + case=case, pair=pair, kind=kind, label=label, enable_preemption=enabled + ) + try: + trial(model, args, requests, args.output, label, state, raw) + except BaseException: + state["error"] = traceback.format_exc() + raw.write(json.dumps(state) + "\n") + raw.flush() + raise + raw.write(json.dumps(state) + "\n") + raw.flush() + rows.append(state) + outputs[enabled] = { + row["request_id"]: (row["token_ids"], row["exit_depths"]) + for row in state["requests"] + } + if outputs[False] != outputs[True]: + raw.write( + json.dumps(dict(case=case, pair=pair, error="paired token/depth mismatch")) + + "\n" + ) + raw.flush() + raise AssertionError("paired traffic changed generated token/depth output") + summary = dict( + cases=paired_summary(rows), + full_matrix={row["case"] for row in rows if row["kind"] == "measurement"} == set(CASES), + default_enable_preemption=False, + profiles_excluded_from_timing=True, + ) + (args.output / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/docs/preemption-benchmarks.md b/docs/preemption-benchmarks.md new file mode 100644 index 00000000..70cd8500 --- /dev/null +++ b/docs/preemption-benchmarks.md @@ -0,0 +1,58 @@ +# Local preemption benchmarks + +`benchmarks.preemption` compares preemption off/on with the same seeded arrivals, +token prompts, fixed output limits, KV budget and priority policy. It covers all +eight combinations of equal/mixed priority, uniform/mixed short-long requests, +and low/high offered load. Higher-priority arrivals use priority 0; lower-priority +requests use 10. The first active wave has lower priority in mixed-priority cells. + +These are controlled scheduling workloads on a real local Ouro checkpoint, +not a language-quality evaluation. Uniform requests use 128 input/64 output +tokens. Mixed lengths use 256/128 and 32/16. Greedy decoding ignores EOS so +both arms do the same requested work. Paired output token IDs and exit depths +must match; all partial records and failures remain in `raw.jsonl`. + +Choose and freeze the two request rates using a separate pilot on the actual +platform. Low/high are offered-load labels: inspect the measured waiting queue +and lateness to confirm that they exercise the intended regimes. Do not adjust +rates after looking at the paired improvement. Run only inside the existing +resource reservation, using an already verified local checkpoint and runtime. + +```bash +python -m benchmarks.preemption --model /path/to/Ouro-1.4B \ + --output results/preemption-start0 --low-rps 0.5 --high-rps 4 \ + --pairs 5 --seed 1729 --profile-case mixed_mixed_high +python -m benchmarks.preemption --model /path/to/Ouro-1.4B \ + --output results/preemption-start1 --low-rps 0.5 --high-rps 4 \ + --pairs 5 --seed 2718 --profile-case mixed_mixed_high +``` + +The example rates are placeholders, not measured platform capacity. Each command +is a fresh process. Use `--decode native` to check ordinary decoding separately, +and `--cuda-graphs` for a separate graph-enabled comparison. A `--plan-only` +invocation imports no Torch or model. `--case` can select a diagnostic subset; +its summary records incomplete matrix coverage. + +Each arm constructs its own engine and warms the same shapes before measurement. +Pair order alternates off/on and on/off. Timed runs retain per-request arrival, +enqueue, first output and completion times; whole-run throughput; peak PyTorch +memory; and preemption/resumption counts. TTFT and E2E include enqueue lateness +relative to the frozen arrival schedule. Enqueue-relative values are reported +separately. Mean, median and p95 metrics are broken down by priority and length; +`summary.json` retains each pair's on/off ratio without discarding slow values. +Model loading, checkpoint hashing and warmup are outside timing. + +Optional profiling runs occur after the timing matrix. They use the repository's +[PyTorch profiling interface](profiling.md), with shapes, stacks and memory +enabled for a finite step window. Capture and artifact-worker completion are +joined before exit. Inspect each profile status and its recorded window: a short +window may miss the actual preemption. Instrumented timing is excluded from the +performance summary. `--profile-only --profile-case mixed_mixed_high` captures +that case without rerunning a completed timing matrix. Profiling artifacts retain host identity and must stay +private when the platform's disclosure policy requires that. + +`SchedulerConfig.enable_preemption` remains `False`. The harness never changes +that default. Enabling it by default needs consistent TTFT benefits with a small +E2E penalty across this matrix and both process starts, including the affected +priority/length groups. CPU regressions and unexecuted benchmark plans provide +no evidence for that change. diff --git a/tests/test_speculative.py b/tests/test_speculative.py index 7dd5ca16..b1e5a5e9 100644 --- a/tests/test_speculative.py +++ b/tests/test_speculative.py @@ -331,7 +331,6 @@ def test_probability_filters_match_expected_and_keep_boundary_ties(): [ dict(cache_config=CacheConfig(layout="shared")), dict(execution_config=ExecutionConfig(async_scheduling=True)), - dict(scheduler_config=SchedulerConfig(enable_preemption=True)), dict(scheduler_config=SchedulerConfig(mode="no_refill")), ], ) @@ -464,3 +463,136 @@ def run(graphs): assert runner.graphs.fallbacks > 0 else: assert any(rows > 2 for rows, _, _, _ in runner.graphs.entries) + + +@pytest.mark.parametrize( + "device,backend,graphs", + [ + ("cpu", "torch", False), + pytest.param("cuda", "triton", False, marks=pytest.mark.gpu), + pytest.param("cuda", "triton", True, marks=pytest.mark.gpu), + ], +) +def test_priority_preempts_only_between_speculative_rounds_and_resumes(device, backend, graphs): + m = ( + model() + if device == "cpu" + else OuroForCausalLM(tiny_ouro_config(hidden_size=256, head_dim=64)).to( + "cuda", torch.bfloat16 + ) + ) + prompts = {"low": [2, 3, 4], "high": [7, 8, 9]} + params = { + "low": SamplingParams(max_tokens=10, ignore_eos=True, priority=10), + "high": SamplingParams(max_tokens=7, ignore_eos=True, priority=0), + } + expected = { + rid: LLM(m, attention_backend=backend).generate([prompt], params[rid])[0] + for rid, prompt in prompts.items() + } + e = engine( + m, + k=3, + cache_config=CacheConfig(64, 16 if graphs else 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=1, + max_num_batched_tokens=4, + prefill_chunk_size=2, + policy="priority", + enable_preemption=True, + ), + attention_backend=backend, + execution_config=ExecutionConfig(cuda_graphs=graphs, cuda_graph_max_batch_size=8), + ) + e.add_request("low", prompts["low"], params["low"]) + while len(e.scheduler.requests["low"].generated_token_ids) < 2: + e.step() + e.add_request("high", prompts["high"], params["high"]) + actual = drain(e) + assert e.preemption.preemptions > 0 + assert e.preemption.resumptions > 0 + assert {rid: out.token_ids for rid, out in actual.items()} == { + rid: out.token_ids for rid, out in expected.items() + } + assert {rid: out.exit_depths for rid, out in actual.items()} == { + rid: out.exit_depths for rid, out in expected.items() + } + assert e.cache_manager.num_used_blocks == 0 + + +@pytest.mark.parametrize("temperature", [0, 0.8]) +def test_speculative_preemption_preserves_request_rng(temperature): + m = model() + params = { + "low": SamplingParams( + max_tokens=12, ignore_eos=True, priority=10, temperature=temperature, top_k=7, seed=42 + ), + "high": SamplingParams( + max_tokens=7, ignore_eos=True, priority=0, temperature=temperature, top_k=7, seed=17 + ), + } + prompts = {"low": [2, 3, 4], "high": [7, 8, 9]} + expected = { + rid: LLM(m, speculative_config=SpeculativeConfig(3)).generate([prompt], params[rid])[0] + for rid, prompt in prompts.items() + } + e = engine( + m, + cache_config=CacheConfig(64, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=1, + max_num_batched_tokens=4, + prefill_chunk_size=2, + policy="priority", + enable_preemption=True, + ), + ) + e.add_request("low", prompts["low"], params["low"]) + while len(e.scheduler.requests["low"].generated_token_ids) < 2: + e.step() + e.add_request("high", prompts["high"], params["high"]) + actual = drain(e) + assert e.preemption.preemptions > 0 and e.preemption.resumptions > 0 + assert {rid: out.token_ids for rid, out in actual.items()} == { + rid: out.token_ids for rid, out in expected.items() + } + assert {rid: out.exit_depths for rid, out in actual.items()} == { + rid: out.exit_depths for rid, out in expected.items() + } + assert not e.preemption.snapshots + assert e.cache_manager.num_used_blocks == 0 + + +def test_abort_suspended_speculative_request_allows_id_reuse(): + m = model() + e = engine( + m, + cache_config=CacheConfig(64, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=1, + max_num_batched_tokens=4, + prefill_chunk_size=2, + policy="priority", + enable_preemption=True, + ), + ) + low = SamplingParams(max_tokens=12, ignore_eos=True, priority=10) + high = SamplingParams(max_tokens=7, ignore_eos=True, priority=0) + e.add_request("low", [2, 3, 4], low) + while len(e.scheduler.requests["low"].generated_token_ids) < 2: + e.step() + suspended = e.scheduler.requests["low"] + e.add_request("high", [7, 8, 9], high) + e.step() + assert "low" in e.preemption.snapshots + assert suspended.stage == Stage.WAITING + assert e.abort_request("low").finished + assert "low" not in e.preemption.snapshots + e.add_request("low", [5, 6], low) + assert e.scheduler.requests["low"] is not suspended + actual = drain(e) + expected = LLM(m, speculative_config=SpeculativeConfig(3)).generate([[5, 6]], low)[0] + assert actual["low"].token_ids == expected.token_ids + assert actual["low"].exit_depths == expected.exit_depths + assert not e.preemption.snapshots + assert e.cache_manager.num_used_blocks == 0 diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 2491dcf0..05184357 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -132,7 +132,7 @@ def __post_init__(self): @dataclass(frozen=True) class SpeculativeConfig: - """Fixed-depth self-speculation; K is explicit until benchmarked.""" + """Fixed-depth self-speculation; set the number of speculative tokens explicitly.""" num_speculative_tokens: int draft_loops: int = 2 diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index c84a1b44..0e7bcb3c 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -237,9 +237,15 @@ def _admit(self): admission increments earlier blocked requests' bypass counts; reaching the limit stops further bypasses so active work can drain. """ + if not self.queues[Stage.WAITING]: + return active_count = sum( r.stage not in (Stage.WAITING, Stage.RECEIVING) for r in self.requests.values() ) + if active_count >= self.config.max_num_seqs and ( + self.config.policy != "priority" or self.preempt_callback is None + ): + return waiting = self._order_waiting_requests() deferred = [] scan_count = min(len(waiting), self.config.admission_scan_limit) diff --git a/vllm_rlt/engine/llm_engine.py b/vllm_rlt/engine/llm_engine.py index a7708960..756fbf0d 100644 --- a/vllm_rlt/engine/llm_engine.py +++ b/vllm_rlt/engine/llm_engine.py @@ -45,10 +45,8 @@ def __init__( raise ValueError("speculative decoding requires fixed-depth ouro exit mode") if self.execution_config.async_scheduling: raise ValueError("speculative decoding currently requires synchronous execution") - if scheduler_config.enable_preemption or scheduler_config.mode != "refill": - raise ValueError( - "speculative decoding requires refill scheduling without preemption" - ) + if scheduler_config.mode != "refill": + raise ValueError("speculative decoding requires refill scheduling") # Retained q distributions and sampling scratch coexist with target # logits. Reserve beyond ordinary prefill/core profiling, even when # requests later choose sampling rather than greedy. @@ -196,6 +194,16 @@ def add_request( def has_unfinished_requests(self) -> bool: return self.scheduler.has_unfinished_requests + def _release_request_state(self, request): + """Release device resources for suspension, retaining the registered request.""" + self.model_runner.release(request.request_id) + self.cache_manager.poll_prefixes() + self.cache_manager.free(request.request_id) + for queue in self.scheduler.queues.values(): + while request.request_id in queue: + queue.remove(request.request_id) + request.hidden_state = request.input_token_tensor = None + def abort_request(self, request_id: str) -> RequestOutput: self.preemption.discard_snapshot(request_id) self.model_runner.release(request_id) diff --git a/vllm_rlt/engine/preemption.py b/vllm_rlt/engine/preemption.py index 7c560117..0839e300 100644 --- a/vllm_rlt/engine/preemption.py +++ b/vllm_rlt/engine/preemption.py @@ -76,6 +76,8 @@ def preempt(self, requester, *, priority_only=False): e = self.engine e.model_runner.synchronize() cache = e.cache_manager + # snapshot copies KV pages individually to CPU, avoiding a temporary + # GPU tensor to gather them while device memory is already under pressure. snapshot = dict( stage=victim.stage, kv=cache.snapshot(victim.request_id), @@ -84,13 +86,7 @@ def preempt(self, requester, *, priority_only=False): if victim.input_token_tensor is None else victim.input_token_tensor.cpu().clone(), ) - e.model_runner.release(victim.request_id) - cache.poll_prefixes() - cache.free(victim.request_id) - for q in e.scheduler.queues.values(): - while victim.request_id in q: - q.remove(victim.request_id) - victim.hidden_state = victim.input_token_tensor = None + e._release_request_state(victim) self.snapshots[victim.request_id] = snapshot e.scheduler.enqueue(victim, Stage.WAITING) self.preemptions += 1