[Feature][MRV2] Adapt extract_hidden_states for Model Runner V2 - #14699
[Feature][MRV2] Adapt extract_hidden_states for Model Runner V2#14699yjyang62 wants to merge 4 commits into
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enables Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
66db379 to
f9b2fc6
Compare
![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
Suggested PR Title:\n\nmarkdown\n[Attention][Feature] Support Model Runner V2 for extract_hidden_states speculative decoding on Ascend\n\n\nSuggested PR Summary:\n\nmarkdown\n### What this PR does / why we need it?\nThis pull request adds support for Model Runner V2 (MRV2) for the `extract_hidden_states` speculative decoding mode on Ascend. It thin-wraps upstream vLLM's `ExtractHiddenStatesSpeculator`, adds dispatching in `init_speculator`, ensures auxiliary hidden state outputs are enabled, and handles allocation and reshaping for `HiddenStateCacheSpec` on a single-tensor path.\n\nFeedback on the changes suggests reusing the existing helper function `_allocate_int8_cache_tensor` in `vllm_ascend/worker/v2/attn_utils.py` to avoid duplicating aligned/unaligned tensor allocation logic.\n\n### Does this PR introduce _any_ user-facing change?\nYes, users can now enable Model Runner V2 for `extract_hidden_states` on Ascend by setting the environment variable `VLLM_USE_V2_MODEL_RUNNER=1`.\n\n### How was this patch tested?\nThe changes were tested with new end-to-end tests in `test_extract_hidden_states.py` covering dense and hybrid models with MRV2, as well as unit tests in `test_attn_utils_v2.py` and `test_extract_hidden_states_speculator_v2.py`.\n
| if vllm_config.kv_transfer_config is None: | ||
| tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device) | ||
| else: | ||
| tensor = torch.zeros( | ||
| kv_cache_tensor.size + alignment, | ||
| dtype=torch.int8, | ||
| device=device, | ||
| ) | ||
| tensor = _align_memory(tensor, alignment)[: kv_cache_tensor.size] | ||
|
|
||
| if has_mamba and has_hidden: | ||
| # Keep Mamba and hidden-state dumps on separate physical buffers | ||
| # so float32 SSM writes cannot corrupt bfloat16 hidden states. | ||
| for layer_name in kv_cache_tensor.shared_by: | ||
| if is_hidden_state_cache_spec(layer_kv_cache_spec[layer_name]): | ||
| if vllm_config.kv_transfer_config is None: | ||
| hidden_tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device) | ||
| else: | ||
| hidden_tensor = torch.zeros( | ||
| kv_cache_tensor.size + alignment, | ||
| dtype=torch.int8, | ||
| device=device, | ||
| ) | ||
| hidden_tensor = _align_memory(hidden_tensor, alignment)[: kv_cache_tensor.size] | ||
| kv_cache_raw_tensors[layer_name] = hidden_tensor | ||
| else: | ||
| kv_cache_raw_tensors[layer_name] = tensor | ||
| else: | ||
| for layer_name in kv_cache_tensor.shared_by: | ||
| kv_cache_raw_tensors[layer_name] = tensor |
There was a problem hiding this comment.
Instead of duplicating the aligned/unaligned tensor allocation logic for tensor and hidden_tensor, you can reuse the existing helper function _allocate_int8_cache_tensor defined in the same file. This improves code maintainability and readability.
tensor = _allocate_int8_cache_tensor(kv_cache_tensor.size, alignment, device)
if has_mamba and has_hidden:
# Keep Mamba and hidden-state dumps on separate physical buffers
# so float32 SSM writes cannot corrupt bfloat16 hidden states.
for layer_name in kv_cache_tensor.shared_by:
if is_hidden_state_cache_spec(layer_kv_cache_spec[layer_name]):
hidden_tensor = _allocate_int8_cache_tensor(kv_cache_tensor.size, alignment, device)
kv_cache_raw_tensors[layer_name] = hidden_tensor
else:
kv_cache_raw_tensors[layer_name] = tensor
else:
for layer_name in kv_cache_tensor.shared_by:
kv_cache_raw_tensors[layer_name] = tensor
github-actions
Bot
added
documentation
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
f9b2fc6 to
10bcc54
Compare
Port upstream vLLM PR #49811 so Ascend NPUModelRunner V2 can dispatch extract_hidden_states, request aux hidden outputs, and keep HiddenStateCacheSpec on a single-tensor allocate/reshape path. Signed-off-by: yjyang62 <yangjinyang5@huawei.com>
Drop the vendored propose/load implementation and thin-wrap upstream vLLM PR #49811 so Ascend only owns dispatch, aux enabling, and NPU KV layout for HiddenStateCacheSpec. Signed-off-by: yjyang62 <yangjinyang5@huawei.com>
10bcc54 to
4bddca1
Compare
yjyang62
changed the title
Signed-off-by: yjyang62 <yangjinyang5@huawei.com>
Upstream vLLM PR #49811 is not in the CI pin yet, so mypy fails with import-not-found on the thin Ascend wrapper. Mark the import until the module lands in the pinned vLLM tree. Signed-off-by: yjyang62 <yangjinyang5@huawei.com> Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
1 participant
What this PR does / why we need it?
Upstream vLLM PR #49811 added Model Runner V2 support for
extract_hidden_states. Ascend already supports this method on Model Runner V1, but Ascend MRV2 previously raisedNotImplementedErrorininit_speculator, and the Ascend V2 KV allocate/reshape path treatedHiddenStateCacheSpeclike MLA K/V (split tensors).This PR adapts Ascend MRV2:
ExtractHiddenStatesSpeculatorasAscendExtractHiddenStatesSpeculatorand dispatches it from Ascendinit_speculator(depends on upstream PR #49811 / a vLLM build that ships that module)use_aux_hidden_state_outputs=Trueforextract_hidden_statesinNPUModelRunnerwhen the pinned vLLM omits the method from the GPU allow-listHiddenStateCacheSpec/cache_only_layerson a single-tensor allocate + reshape path in MRV2attn_utilsVLLM_USE_V2_MODEL_RUNNER=1Does this PR introduce any user-facing change?
Yes. With
VLLM_USE_V2_MODEL_RUNNER=1, users can runextract_hidden_stateson Model Runner V2 (once the Ascend image/vLLM includes PR #49811).How was this patch tested?
tests/ut/worker/test_extract_hidden_states_speculator_v2.pytests/ut/worker/test_attn_utils_v2.py::test_mrv2_allocates_and_reshapes_hidden_state_cachetests/e2e/pull_request/one_card/spec_decode/test_extract_hidden_states.py(dense_eager_mrv2,hybrid_dummy_eager_mrv2). NPU e2e was not run in the cloud agent environment (no Ascend device); please verify on Atlas hardware.pytest -sv tests/ut/worker/test_extract_hidden_states_speculator_v2.py \ tests/ut/worker/test_attn_utils_v2.py::test_mrv2_allocates_and_reshapes_hidden_state_cache export VLLM_USE_V2_MODEL_RUNNER=1 pytest -sv tests/e2e/pull_request/one_card/spec_decode/test_extract_hidden_states.py -k mrv2 - vLLM version: v0.27.1 - vLLM main: https://github.com/vllm-project/vllm/commit/58d3918e3ea0a544ffedadad2ba84559e9c51d8f