Spec decode warmup support

[jerrychenhf]

GAUDISW-242931

Because currently spec decode flatten the spec decode tokens into [batch_size * num_tokens, 1], we can warmup the decode shapes as it was. The thing changed is the maximum batch_size we should warmup in the configuration because the real batch size is batch_size * num_tokens which is num_tokens (1 + num_speculative_tokens) times of original batch size.

The thing to care in the warmup is the draft token (and block) space for the proposing process in eagle. We need to leave out the num_speculative_tokens space to use by propose for eagle.

Other care needs to be taken (already done in the PR of support num_speculative_tokens > 1) is warmup will be run in compile only mode without the real computation happening. So the operations for prepare_attn_metadata in the drafter which depends on the real position values must be done on CPU)

Another issue of handling no spec decode tokens for decode phase has already been handled #593

[github-actions]

🚧 CI Blocked

The main CI workflow was not started for the following reason:

Your branch is behind the base branch. Please merge or rebase to get the latest changes.

[github-actions]

✅ CI Passed

All checks passed successfully against the following vllm commit:
0353d2e162cbda776d9dbfe026e65303204a7f1f

[github-actions]

✅ CI Passed

All checks passed successfully against the following vllm commit:
0353d2e162cbda776d9dbfe026e65303204a7f1f

[jerrychenhf]

@xuechendi Would you please have a review on this?

[xuechendi]

same question here

[xuechendi]

Good catch!

[jerrychenhf]

Originally, I take the simplified approach that it's the user to configure the bucketing parameters properly to cover the shapes (batch_size, 1, context_blocks) for spec decode.
The new commit provide the solution in bucket manager to automatically generate possible new buckets for spec decode based on the buckets configured by user (seed buckets). This somewhat simplify the user configuration.

Please note: Current spec decode has some tricky things as to the bucketed batch size and context blocks:

1. We first get a padded_batch_size based on the num_decodes:

padded_batch_size = self.bucketing_manager.find_decode_bucket(num_decodes, sum(num_blocks))[0]

2. And then a virtual batch size is got from padded_batch_size * num_tokens. This is the real running batch size, which is not ganrateed to be covered by any bucketing (if there is no special handling).
3. And then bucketed context_blocks is find using the virtual batch size (but we used only the bucketed blocks but the bucketed batch size is not used, the virtual batch size = padded_batch_size * num_tokens is still used).

block_bucket_size = \
                self.bucketing_manager.find_decode_bucket(batch_size,
                                                          len(block_list))[2]

[github-actions]

🚧 CI Blocked

The main CI workflow was not started for the following reason:

Your branch is behind the base branch. Please merge or rebase to get the latest changes.

[xuechendi]

@jerrychenhf , please add profiling to ticket to indicate the hpu graph have been captured for spec decode.

[jerrychenhf]

Updated the test results to GAUDISW-242931. With warmup enabled with this PR, running "PT_HPU_LAZY_MODE=1 python tests/full_tests/spec_decode.py --task eagle3 --batch_size 8 --osl 2048--num_spec_tokens 2", no configurations were shown as "not warmup up".

[github-actions]

✅ CI Passed

All checks passed successfully against the following vllm commit:
d7284a2604ef3fe96f0779309caafb59860704bb