Store internal states of speculative decode

[vladislavkruglikov]

I want to get internal states of speculative decode such as history of proposals and scores in the response as well as completion

To achieve my goal I generally need to get proposals and proposal scores from model executor

vllm/vllm/engine/async_llm_engine.py

Lines 282 to 283 in d3d9cb6

	output = await self.model_executor.execute_model_async(
	execute_model_req)

After this pass to function that is responsible for update of internal state of sequence with outputs from model executor

vllm/vllm/engine/async_llm_engine.py

Lines 287 to 289 in d3d9cb6

	request_outputs = self._process_model_outputs(
	output, scheduler_outputs.scheduled_seq_groups,
	scheduler_outputs.ignored_seq_groups, seq_group_metadata_list)

But the problem that I face is that list[SamplerOutput] that is being returned from the model executor and speculative worker under the hood stores data about each token not a whole sequence. I mean I could possibly append proposals and proposal scores for every token in sampler output but this would be just weird and scary so I need to figure out some prettier way to do that

One way is to enrich return type of worker to allow it store some metadata except sampler output for example something like this would work

class WorkerResponse:
    sampler_outputs: list[SamplerOutput]

class SpeculativeWorkerResponse(WorkerResponse):
    proposals: SpeculativeProposals
    proposal_scores: SpeculativeScores

So change from this

vllm/vllm/spec_decode/spec_decode_worker.py

Lines 333 to 337 in d3d9cb6

	@torch.inference_mode()
	def execute_model(
	self,
	execute_model_req: Optional[ExecuteModelRequest] = None
	) -> List[SamplerOutput]:

To this

@torch.inference_mode()
def execute_model(
    self,
    execute_model_req: Optional[ExecuteModelRequest] = None
) -> SpeculativeWorkerResponse:

This data structure can be returned from speculative worker and after all returned by model executor and passed to post process function that would append proposals and proposal scores to specific sequence data items

My question is what do you think about extending response type of all workers to allow it to have additional data or do you have your own thoughts on this?

[vladislavkruglikov]

@WoosukKwon @Yard1 @youkaichao @cadedaniel what are your thoughts on this?

[cadedaniel]

What's the motivation to add this to the live-serving scenario? Unfortunately we are very latency constrained, so additional features must be justified, otherwise we will dismiss the latency reduction in live serving.

[cadedaniel]

BTW, similar discussion happening here #7586

[vladislavkruglikov]

Metric is not the main purpose of this feature. Since our algorithm uses previous draft in order to predict next draft we need this feature. You can not implement this algorithm and do not store the previous drafts

What if I will add an additional flag such as return_speculative_decode_states which will add SpeculativeProposerSamplerOutput and SpeculativeScorerSamplerOutput to the output and engine will check if that flag is set in order to start processing it. This will definitely help to remove GPU synchronization if you do not want to use additional features you do not want to pay for

[cadedaniel]

Metric is not the main purpose of this feature. Since our algorithm uses previous draft in order to predict next draft we need this feature. You can not implement this algorithm and do not store the previous drafts

I am confused; why do you want to return these previous draft states to the user? If you want to improve your drafting live, e.g. https://arxiv.org/abs/2310.07177, you want to keep these in the SpecDecodeWorker.

[vladislavkruglikov]

If we would store speculative decode states in speculative decode worker this would definitely solve our problem. But one more thing that we need it so have ability to get internal states of speculative decoder as fast as possible. If we would store internal states as logs then it would be much slower to access it rather then directly receiving them as API response

[cadedaniel]

If we would store speculative decode states in speculative decode worker this would definitely solve our problem

Great. Feel free to make a PR for this, or an RFC proposing the change. In general, if we can keep speculative decoding features inside the SpecDecodeWorker, it will be easier to merge.

But one more thing that we need it so have ability to get internal states of speculative decoder as fast as possible. If we would store internal states as logs then it would be much slower to access it rather then directly receiving them as API response

Can you help me understand the use case better? Who is the user, why does the user need it, how fast do you need it, why can't it stay inside the SpecDecodeWorker.