Conversation
The worker-side MTP transfer (vllm-project#13454) only updated KVPoolWorker.num_layers to include MTP draft layers; KVPoolScheduler.num_layers stayed at model_config.get_num_layers(), which excludes MTP. This broke the 'Keep this in sync with pool_worker.py' invariant and left the scheduler's per-layer lookup keys and GVA allocation sizing short of the MTP layer, causing fluctuating MTP draft acceptance under memcache/gva layerwise. Add a shared extract_physical_layer_index helper in config_data (stdlib re, mirrors the worker). KVPoolScheduler now derives num_layers (and keys_per_block_hash) from kv_cache_config.kv_cache_groups[*].layer_names so MTP layers are counted; for PP the scheduler receives kv_cache_configs[0] (PP rank 0, no MTP), matching its [self.pp_rank] query scope. KVPoolWorker._extract_physical_layer_index delegates to the shared helper (behavior-preserving; drops the inline regex dependency). Signed-off-by: tyy0829 <1455207791@qq.com>
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 addresses a synchronization issue between the KV pool scheduler and worker regarding MTP (Multi-Token Prediction) draft layers. Previously, the scheduler's layer count excluded MTP layers, leading to inconsistent GVA allocation and potential failures in draft acceptance. By unifying the layer indexing logic and aligning the scheduler's layer accounting with the worker's authoritative configuration, this change ensures consistent behavior across the distributed KV cache system. 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
|
|
👋 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. Tip 💡 Consider Linking a Related Issue or RFCYour PR title contains the [BugFix] tag, indicating a bug fix or new feature. Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:
🙏 Thanks for helping us keep the project well-organized! |
![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:
[Ops][Feature] Support MTP layers in KVPoolScheduler and share physical layer index extractionSuggested PR Summary:
### What this PR does / why we need it?
This PR introduces a shared helper function `extract_physical_layer_index` to extract physical layer indices (including MTP draft layers) from registered KV cache layer names. It updates `KVPoolScheduler` to compute the registered number of layers dynamically from `KVCacheConfig`, ensuring the scheduler's per-layer key generation and GVA allocation stay in sync with `KVPoolWorker`.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added unit tests in `test_pool_scheduler.py` to verify the physical layer index extraction and the scheduler's layer count computation.Review Feedback:
All review comments are valid and point out an actual improvement opportunity: the fallback parameter in extract_physical_layer_index is unused and should be removed to simplify the API.
| def extract_physical_layer_index( | ||
| layer_name: str, num_hidden_layers: int, fallback: int | ||
| ) -> int: | ||
| """Extract the physical layer index from a registered KV cache layer name. | ||
|
|
||
| MTP draft layers are mapped after the main model layers so they sort | ||
| behind the standard attention layers: | ||
|
|
||
| * DeepSeek-style MTP reuses the ``layers.{N}`` naming with | ||
| ``N >= num_hidden_layers`` (e.g. ``model.layers.61.self_attn.attn``). | ||
| * Other models use an ``mtp.{N}`` prefix (e.g. ``model.mtp.0.self_attn``) | ||
| which is mapped to ``num_hidden_layers + N``. | ||
|
|
||
| Shared by ``KVPoolWorker`` and ``KVPoolScheduler`` so both sides derive an | ||
| identical layer layout (including MTP) from the registered layer names. | ||
|
|
||
| Args: | ||
| layer_name: registered KV cache layer name. | ||
| num_hidden_layers: main model hidden layer count, used to offset MTP | ||
| layers that use the ``mtp.{N}`` prefix. | ||
| fallback: value used when an ``mtp.{N}`` prefix is matched but no main | ||
| layer count is available. | ||
| """ |
There was a problem hiding this comment.
The fallback parameter is defined in the function signature and documented in the docstring, but it is completely unused within the function body. Since the fallback logic is already handled by the callers before invoking this function, we should remove the fallback parameter to simplify the API and avoid confusion.
def extract_physical_layer_index(
layer_name: str, num_hidden_layers: int
) -> int:
"""Extract the physical layer index from a registered KV cache layer name.
MTP draft layers are mapped after the main model layers so they sort
behind the standard attention layers:
* DeepSeek-style MTP reuses the ``layers.{N}`` naming with
``N >= num_hidden_layers`` (e.g. ``model.layers.61.self_attn.attn``).
* Other models use an ``mtp.{N}`` prefix (e.g. ``model.mtp.0.self_attn``)
which is mapped to ``num_hidden_layers + N``.
Shared by ``KVPoolWorker`` and ``KVPoolScheduler`` so both sides derive an
identical layer layout (including MTP) from the registered layer names.
Args:
layer_name: registered KV cache layer name.
num_hidden_layers: main model hidden layer count, used to offset MTP
layers that use the ``mtp.{N}`` prefix.
"""| physical_layers.add( | ||
| extract_physical_layer_index(layer_name, num_hidden_layers, fallback) | ||
| ) |
There was a problem hiding this comment.
Remove the unused fallback argument when calling extract_physical_layer_index.
| physical_layers.add( | |
| extract_physical_layer_index(layer_name, num_hidden_layers, fallback) | |
| ) | |
| physical_layers.add( | |
| extract_physical_layer_index(layer_name, num_hidden_layers) | |
| ) |
| m = re.search(r"(\d+)", layer_name) | ||
| return int(m.group(1)) if m else 0 | ||
| num_hidden_layers = getattr(self.hf_config, "num_hidden_layers", self.num_layers) | ||
| return extract_physical_layer_index(layer_name, num_hidden_layers, self.num_layers) |
There was a problem hiding this comment.
Remove the unused self.num_layers fallback argument when calling extract_physical_layer_index.
| return extract_physical_layer_index(layer_name, num_hidden_layers, self.num_layers) | |
| return extract_physical_layer_index(layer_name, num_hidden_layers) |
| self.assertEqual( | ||
| extract_physical_layer_index("model.layers.0.self_attn.attn", 2, 2), 0 | ||
| ) | ||
| self.assertEqual( | ||
| extract_physical_layer_index("model.layers.60.self_attn.attn", 61, 61), 60 | ||
| ) | ||
| # First MTP layer for a 61-layer model (index == num_hidden_layers). | ||
| self.assertEqual( | ||
| extract_physical_layer_index("model.layers.61.self_attn.attn", 61, 61), 61 | ||
| ) | ||
|
|
||
| def test_mtp_prefix_maps_after_main_layers(self): | ||
| self.assertEqual( | ||
| extract_physical_layer_index("model.mtp.0.self_attn", 2, 2), 2 | ||
| ) | ||
| self.assertEqual( | ||
| extract_physical_layer_index("model.mtp.1.self_attn", 2, 2), 3 | ||
| ) | ||
|
|
||
| def test_fallback_when_no_numeric_token(self): | ||
| self.assertEqual(extract_physical_layer_index("no_index", 2, 2), 0) |
There was a problem hiding this comment.
Update the unit tests to match the simplified signature of extract_physical_layer_index by removing the unused third argument.
self.assertEqual(
extract_physical_layer_index("model.layers.0.self_attn.attn", 2), 0
)
self.assertEqual(
extract_physical_layer_index("model.layers.60.self_attn.attn", 61), 60
)
# First MTP layer for a 61-layer model (index == num_hidden_layers).
self.assertEqual(
extract_physical_layer_index("model.layers.61.self_attn.attn", 61), 61
)
def test_mtp_prefix_maps_after_main_layers(self):
self.assertEqual(
extract_physical_layer_index("model.mtp.0.self_attn", 2), 2
)
self.assertEqual(
extract_physical_layer_index("model.mtp.1.self_attn", 2), 3
)
def test_fallback_when_no_numeric_token(self):
self.assertEqual(extract_physical_layer_index("no_index", 2), 0)
|
Does the main branch not have this issue? |
|
Was this issue not discovered when testing #13454? |
2 participants
What this PR does / why we need it?
Follow-up to #13454, which only updated
KVPoolWorker.num_layersto include MTP draft layers.KVPoolScheduler.num_layerswas left atmodel_config.get_num_layers(), which excludes MTP -- breaking the "Keep this in sync with pool_worker.py" invariant (pool_scheduler.py:179) and leaving the scheduler's per-layer lookup keys / GVA allocation sizing (num_layer_keys,keys_per_block_hash,split_layers) short of the MTP layer. Under memcache/gva layerwise this caused fluctuating MTP draft acceptance: the scheduler declared a block "hit" without verifying the MTP layer key, then the worker attempted an MTP load that could fail depending on whether the producer had saved the MTP KV yet.Changes:
config_data.py: add sharedextract_physical_layer_indexhelper (stdlibre; mirrorsKVPoolWorker._extract_physical_layer_index). Covers both DeepSeek-stylemodel.layers.{N>=num_hidden_layers}andmodel.mtp.{N}naming.pool_scheduler.py:KVPoolSchedulernow derivesnum_layers(and recomputeskeys_per_block_hash) fromkv_cache_config.kv_cache_groups[*].layer_names, so MTP layers are counted. For PP, the scheduler receiveskv_cache_configs[0](PP rank 0, no MTP), matching its[self.pp_rank]query scope.pool_worker.py:_extract_physical_layer_indexdelegates to the shared helper (behavior-preserving; drops the inlineregexdependency).tests/ut/distributed/ascend_store/test_pool_scheduler.py: add tests for DeepSeek-style MTP,mtp..-prefix MTP, no-config fallback, andkeys_per_block_hashunder gva layerwise, plus directextract_physical_layer_indexunit tests.Refs #13454.
Does this PR introduce any user-facing change?
No user-facing API changes. Internal fix for KV pool scheduler layer accounting including MTP layers.
How was this patch tested?
Unit tests added/updated:
tests/ut/distributed/ascend_store/test_pool_scheduler.pytest_num_layers_includes_deepseek_mtp_layertest_num_layers_includes_mtp_prefix_layertest_num_layers_without_kv_cache_config_falls_backtest_keys_per_block_hash_includes_mtp_for_gvaTestExtractPhysicalLayerIndex(helper unit tests)Manual verification (standalone script with project
_mock_depsshims): helper extraction, worker delegation, schedulernum_layers/keys_per_block_hashfor MTP/gva all pass;py_compileclean.CI verification (please run
pytest tests/ut/distributed/ascend_store/test_pool_scheduler.py tests/ut/distributed/ascend_store/test_pool_worker.py+bash format.sh ci).vLLM version: v0.23.0
vLLM main: vllm-project/vllm@ee0da84