[Speculative decoding 4/9] Lookahead scheduling for speculative decoding

[cadedaniel]

This PR introduces the concept of lookahead scheduling. Lookahead scheduling is where we allocate KV slots for each sequence in a decode batch that do not have any token assigned ("empty slots"). Speculative decoding fills these KV slots with the KV of speculative tokens when running the target model. Furthermore, speculative decoding involving a proposal method that has KV cache will also use the KV slots in normal autoregressive generation.

See these step-by-step examples explaining how Lookahead scheduling works.

Note: we could use scratch space for these KV, however in the case where tokens are accepted we would need to copy the accepted KV from the scratch space to the allocated KV slots. By allocating them ahead-of-time, we save us the complexity of scheduling such a memcpy.

Testing

• This PR finishes the copy-on-write scheduler integration from [Core] [Bugfix] Refactor block manager subsystem for better testability #3492 (since now there can be multiple CoW per append_slots).
  • A test is added where we verify v1/v2 block manager output equality when there are copy-on-writes: test_v1_v2_greedy_equality_with_cow
• test_lookahead_greedy_equality_with_preemption tests equality of generation when lookahead is enabled vs disabled, and includes preemption

Temporary flag

A temporary flag --num-lookahead-slots is added to facilitate testing. It will be removed in PR 6/9 of the speculative decoding oss plan.

[cadedaniel]

Ready for review cc @LiuXiaoxuanPKU

[LiuXiaoxuanPKU]

Thanks for the PR! Just left some questions.

[LiuXiaoxuanPKU]

what's the definition of num_new_tokens here?

[cadedaniel]

This is the number of tokens that do not have a slot allocated in the block table. Will add a comment.

[LiuXiaoxuanPKU]

From my understanding, the num_new_tokens == len(unseen_token_ids)?

[cadedaniel]

yes. I will rename them so it's consistent.

[cadedaniel]

added

[LiuXiaoxuanPKU]

num_lookahead_slots is not used in this function, is it expetced?

[cadedaniel]

Not expected, will add ensure_num_empty_slots call with num_lookahead_slots + see why the test didn't catch this.

[cadedaniel]

fixed; added test.

[LiuXiaoxuanPKU]

Nit: I just feel the two functions below _can_append_slots and _can_swap_in are a bit shallow and do not hide much complexity. Maybe we can call them directly above?

[cadedaniel]

pr feedback applied @LiuXiaoxuanPKU

[cadedaniel]

added

[cadedaniel]

fixed; added test.

[animan42]

@cadedaniel really awesome series of changes! I assume the answer is no, but does the draft model also have it's own KV cache? If yes, where is it created and updated?

[cadedaniel]

@cadedaniel really awesome series of changes! I assume the answer is no, but does the draft model also have it's own KV cache? If yes, where is it created and updated?

Thanks! Great question. The draft model can have KV, and by default it does. The spec decode worker has a proposer worker which can be a normal vLLM Worker, with KV cache and everything. The number of blocks in the draft KV and target KV is calculated by this function:

vllm/vllm/spec_decode/spec_decode_worker.py

Lines 97 to 116 in b3104b2

	def determine_num_available_blocks(self) -> Tuple[int, int]:
	"""Determine the number of cache blocks to use.
	This is done by profiling the scorer model (which is typically the
	larger of the two). Then the total memory which would be used by the
	scorer cache is divided evenly between the proposer and scorer model KV,
	such that the number of blocks is equal in both KV caches.
	"""
	num_gpu_blocks, num_cpu_blocks = (
	self.scorer_worker.determine_num_available_blocks())
	scorer_cache_block_size_bytes = (
	self.scorer_worker.get_cache_block_size_bytes())
	proposer_cache_block_size_bytes = (
	self.proposer_worker.get_cache_block_size_bytes())
	new_num_gpu_blocks = split_num_cache_blocks_evenly(
	scorer_cache_block_size_bytes, proposer_cache_block_size_bytes,
	num_gpu_blocks)
	return new_num_gpu_blocks, num_cpu_blocks

As for the mapping between token ids and block indices, we keep things simple by using the same block mapping for the draft and target models. We can do this because of the code linked above, which ensures that the draft and target KV have the same amount of logical KV space. With this, the KV of the draft model is populated in lockstep with the target model.

There are more details, like how proposal KV are handled, but this is the gist!

[xunfeng1980]

  File "/data/anaconda3/envs/qwen-q/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/data/anaconda3/envs/qwen-q/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/data/vllm/vllm/entrypoints/openai/api_server.py", line 157, in <module>
    engine = AsyncLLMEngine.from_engine_args(
  File "/data/vllm/vllm/engine/async_llm_engine.py", line 347, in from_engine_args
    engine = cls(
  File "/data/vllm/vllm/engine/async_llm_engine.py", line 311, in __init__
    self.engine = self._init_engine(*args, **kwargs)
  File "/data/vllm/vllm/engine/async_llm_engine.py", line 421, in _init_engine
    return engine_class(*args, **kwargs)
  File "/data/vllm/vllm/engine/llm_engine.py", line 119, in __init__
    self.model_executor = executor_class(
  File "/data/vllm/vllm/executor/gpu_executor.py", line 37, in __init__
    assert (not speculative_config
AssertionError: Speculative decoding not yet supported for GPU backend