Repository navigation
Conversation
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request generalizes the identification of EAGLE and MTP draft attention layers by introducing an is_eagle flag in the attention specification. This replaces model-specific heuristics in the KV cache grouping logic, specifically for DeepSeek-V4, with a more robust mechanism that propagates the flag from the model layers to the KVCacheCoordinator. I have no feedback to provide as there were no review comments to assess.
|
@ivanium @heheda12345 Would you mind taking a look? Thanks! |
|
|
||
|
|
||
| def _annotate_eagle_groups_deepseek_v4( | ||
| def _annotate_eagle_groups( |
There was a problem hiding this comment.
Ideally we can annotate kv_cache_group.is_eagle when building KVCacheGroupSpec and avoid this function. Would that be too intrusive?
There was a problem hiding this comment.
I think it would be a bit intrusive. And it also introduced some duplicate code to judge is_eagle_group. get_kv_cache_groups playing a role as orchestrator to do it once looks reasonable. I created a commit 14d5684 here.
There was a problem hiding this comment.
We can simplify them to 6 lines inline code in get_kv_cache_groups.
spec_config = vllm_config.speculative_config
if spec_config is not None and spec_config.use_eagle():
eagle_layer_names = {n for n, s in kv_cache_spec.items() if s.is_eagle}
for group in kv_cache_groups:
if eagle_layer_names & set(group.layer_names):
group.is_eagle_group = True
|
Thanks for the effort! I left one comment, and some more general ones:
|
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Thanks for reviewing! I have moved |
|
|
||
| return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] | ||
| layer_names = list(spec.kv_cache_specs.keys()) | ||
| is_eagle_group = bool(eagle_layer_names and eagle_layer_names & set(layer_names)) |
There was a problem hiding this comment.
I prefer a more intuitive code like this:
| is_eagle_group = bool(eagle_layer_names and eagle_layer_names & set(layer_names)) | |
| is_eagle_group = bool(eagle_layer_names) and bool( | |
| eagle_layer_names.intersection(current_group_layer_names) | |
| ) |
| full_mla_group = KVCacheGroupSpec( | ||
| layer_names=list(full_mla_spec.kv_cache_specs.keys()), | ||
| kv_cache_spec=full_mla_spec, | ||
| is_eagle_group=bool( |
| layer_names=group_layer_names, | ||
| kv_cache_spec=sub_sm_spec, | ||
| is_eagle_group=bool( | ||
| eagle_layer_names and eagle_layer_names & set(group_layer_names) |
There was a problem hiding this comment.
ditto. and let's extract a function for determining whether it is a eagle group.
| # time instead of through a separate post-hoc pass. | ||
| spec_config = vllm_config.speculative_config | ||
| if spec_config is not None and spec_config.use_eagle(): | ||
| eagle_layer_names = {n for n, s in kv_cache_spec.items() if s.is_eagle} |
There was a problem hiding this comment.
Let's clarify the meaning of the name more clearly.
| eagle_layer_names = {n for n, s in kv_cache_spec.items() if s.is_eagle} | |
| eagle_layer_names = {layer_name for layer_name, spec in kv_cache_spec.items() if spec.is_eagle} |
| cache_dtype_str=self.cache_config.cache_dtype, | ||
| alignment=576, # NOTE: FlashMLA requires 576B alignment | ||
| model_version="deepseek_v4", | ||
| is_eagle=self.is_eagle, |
| dtype=torch.uint8, | ||
| prefix=f"{prefix}.swa_cache", | ||
| cache_config=cache_config, | ||
| is_eagle=is_eagle, |
There was a problem hiding this comment.
Instead of passing by is_eagle through a long link, we can get the info about whether it is eagle layer in this way:
is_eagle= True if extract_layer_index(prefix) > vllm_config.model_config.hf_config.num_hidden_layers else False
|
This pull request has merge conflicts that must be resolved before it can be |
|
This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you! |
Purpose
Follow #39690.
Currently,
_annotate_eagle_groups_deepseek_v4is model-specific method. Other models would all go into the fallback logic inside. This PR is introducingis_eaglein AttentionSpec to make_annotate_eagle_groupsgeneral for all models.Test Plan
Gemma4
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.