Repository navigation
[BugFix] [v0.23.0]Fix invalid GVA handling in layerwise KV pool - #12643
Conversation
- _alloc_gvas_for_save: only cache and track gva>0 keys, skip failed batch_alloc - _prepare_load_gvas: only call batch_add_lease for keys with valid size (>0) - _prepare_load_gvas: report size<=0 and lease failure blocks to _invalid_block_ids - _prepare_load_gvas: skip invalid_block_ids for multi-group (hybrid) models - build_shared: filter out gva<=0 blocks from block_ids/block_gvas before batch_copy Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.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 issues with invalid Global Virtual Address (GVA) handling in the layerwise KV pool. By ensuring that only valid keys are cached, tracked, and leased, the changes prevent downstream errors during batch operations. Additionally, the PR introduces a safety mechanism to report invalid blocks for recomputation in single-group models while safeguarding hybrid models against potential scheduler inconsistencies. 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. |
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Ops][BugFix] Handle invalid GVAs and prevent scheduler crashes during KV transferSuggested PR Summary:
### What this PR does / why we need it?
This PR improves the robustness of the KV transfer pool by properly handling invalid (zero or negative) Global Virtual Addresses (GVAs) during save and load operations. Specifically, it:
- Filters out invalid GVAs in `build_shared` and logs warnings for zero/negative `base_gvas` in `_build_transfer_arrays`.
- Only allocates and saves keys with valid GVAs in `_alloc_gvas_for_save`.
- Refactors `_prepare_load_gvas` to only lease valid keys and safely report invalid blocks to the scheduler for recomputation (restricted to single-group models to prevent hybrid model scheduler crashes).
Feedback from the review highlights potential issues including:
- Missing filtering of invalid GVAs in `_build_transfer_arrays` which could lead to errors in `batch_copy`.
- Potential `None` dereferences for `ki` and `block_ids_by_group` in `_prepare_load_gvas`.
- Missing defensive check for `None` on `batch_add_lease` results.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Not specified in the PR, but should be verified with unit tests for KV transfer under invalid GVA scenarios.| if base_gvas_arr.size > 0 and np.any(base_gvas_arr <= 0): | ||
| zero_count = int(np.sum(base_gvas_arr <= 0)) | ||
| logger.warning( | ||
| "[KVPOOL] build_transfer layer=%d detected %d zero/negative base_gvas " | ||
| "(base_gvas_sample=%s); these blocks will be skipped in batch_copy", | ||
| layer_id, | ||
| zero_count, | ||
| base_gvas_arr[:5].tolist(), | ||
| ) |
There was a problem hiding this comment.
The warning message states that blocks with zero or negative base_gvas will be skipped in batch_copy. However, the code only logs a warning and does not actually filter them out from base_gvas_arr and block_ids_arr. If invalid GVAs reach this point, they will still be processed and passed to batch_copy, causing errors. We should explicitly filter them out using a boolean mask.
if base_gvas_arr.size > 0 and np.any(base_gvas_arr <= 0):
zero_count = int(np.sum(base_gvas_arr <= 0))
logger.warning(
"[KVPOOL] build_transfer layer=%d detected %d zero/negative base_gvas "
"(base_gvas_sample=%s); these blocks will be skipped in batch_copy",
layer_id,
zero_count,
base_gvas_arr[:5].tolist(),
)
valid_mask = base_gvas_arr > 0
base_gvas_arr = base_gvas_arr[valid_mask]
block_ids_arr = block_ids_arr[valid_mask]| for idx, (ki, key) in enumerate(zip(key_infos, keys)): | ||
| sizes = ki.size() | ||
| gvas.append(ki.gva_list()[0] if sizes and sizes > 0 else 0) | ||
| all_group_load_keys.extend(keys) | ||
| gva = ki.gva_list()[0] if sizes and sizes > 0 else 0 | ||
| gvas.append(gva) | ||
| if gva > 0: | ||
| valid_keys_for_lease.append(key) | ||
| block_idx = load_start_block + idx | ||
| if block_idx < len(block_ids_by_group): | ||
| valid_block_ids.append(int(block_ids_by_group[block_idx])) | ||
| else: | ||
| block_idx = load_start_block + idx | ||
| if block_idx < len(block_ids_by_group): | ||
| invalid_block_ids.append(int(block_ids_by_group[block_idx])) | ||
| logger.warning( | ||
| "load_gvas: req=%s group=%d got invalid gva=%d (size=%d), " | ||
| "block_id=%s load failed", | ||
| request.req_id, group_id, gva, | ||
| sizes if sizes else 0, | ||
| int(block_ids_by_group[block_idx]) if block_idx < len(block_ids_by_group) else "N/A", | ||
| ) |
There was a problem hiding this comment.
There are several potential None dereferences and missing checks here:
ki(returned frombatch_get_key_info) can beNoneif the key is not found or the backend fails, which would cause a crash onki.size().block_ids_by_groupcan beNone(as handled defensively elsewhere in this method), but herelen(block_ids_by_group)andblock_ids_by_group[block_idx]are accessed without checking forNone, which will raise aTypeError.
We should add defensive checks to prevent these crashes.
for idx, (ki, key) in enumerate(zip(key_infos, keys)):
sizes = ki.size() if ki is not None else 0
gva = ki.gva_list()[0] if ki is not None and sizes and sizes > 0 else 0
gvas.append(gva)
block_idx = load_start_block + idx
has_block = block_ids_by_group is not None and block_idx < len(block_ids_by_group)
block_id = int(block_ids_by_group[block_idx]) if has_block else None
if gva > 0:
valid_keys_for_lease.append(key)
if has_block and block_id is not None:
valid_block_ids.append(block_id)
else:
if has_block and block_id is not None:
invalid_block_ids.append(block_id)
logger.warning(
"load_gvas: req=%s group=%d got invalid gva=%d (size=%d), "
"block_id=%s load failed",
request.req_id, group_id, gva,
sizes if sizes else 0,
block_id if has_block else "N/A",
)| if valid_keys_for_lease: | ||
| lease_results = self.m_store.batch_add_lease( | ||
| valid_keys_for_lease, LAYERWISE_READ_LEASE_TTL_MS | ||
| ) | ||
| # Report lease failures as invalid blocks | ||
| for i, lease_res in enumerate(lease_results): | ||
| if lease_res != 0 and i < len(valid_block_ids): | ||
| invalid_block_ids.append(valid_block_ids[i]) | ||
| logger.warning( | ||
| "load_gvas: req=%s group=%d lease failed result=%d, " | ||
| "block_id=%d load failed", | ||
| request.req_id, group_id, lease_res, valid_block_ids[i], | ||
| ) |
There was a problem hiding this comment.
If batch_add_lease fails or returns None due to a backend issue, iterating over lease_results using enumerate will raise a TypeError. We should defensively check if lease_results is not None before iterating.
| if valid_keys_for_lease: | |
| lease_results = self.m_store.batch_add_lease( | |
| valid_keys_for_lease, LAYERWISE_READ_LEASE_TTL_MS | |
| ) | |
| # Report lease failures as invalid blocks | |
| for i, lease_res in enumerate(lease_results): | |
| if lease_res != 0 and i < len(valid_block_ids): | |
| invalid_block_ids.append(valid_block_ids[i]) | |
| logger.warning( | |
| "load_gvas: req=%s group=%d lease failed result=%d, " | |
| "block_id=%d load failed", | |
| request.req_id, group_id, lease_res, valid_block_ids[i], | |
| ) | |
| if valid_keys_for_lease: | |
| lease_results = self.m_store.batch_add_lease( | |
| valid_keys_for_lease, LAYERWISE_READ_LEASE_TTL_MS | |
| ) | |
| if lease_results is not None: | |
| # Report lease failures as invalid blocks | |
| for i, lease_res in enumerate(lease_results): | |
| if lease_res != 0 and i < len(valid_block_ids): | |
| invalid_block_ids.append(valid_block_ids[i]) | |
| logger.warning( | |
| "load_gvas: req=%s group=%d lease failed result=%d, " | |
| "block_id=%d load failed", | |
| request.req_id, group_id, lease_res, valid_block_ids[i], | |
| ) |
3ee4fd5 to
e3a3ea1
Compare
b814edc to
a80bd1c
Compare
The store is already initialized eagerly during MemcacheBackend construction in _init_backend. Calling init_store() again in register_kv_caches is redundant; on the lazy path it re-triggers _setup_store()/SmemBmCreate2, which synchronously allocates the HBM/DRAM memory pool at startup and causes service-launch hangs on old HDK where memory allocation blocks. Signed-off-by: tyy0829 <tyy0829@users.noreply.github.com>
- _alloc_gvas_for_save: only cache and track gva>0 keys, skip failed batch_alloc - _prepare_load_gvas: only call batch_add_lease for keys with valid size (>0) - _prepare_load_gvas: report size<=0 and lease failure blocks to _invalid_block_ids - _prepare_load_gvas: skip invalid_block_ids for multi-group (hybrid) models - build_shared: filter out gva<=0 blocks from block_ids/block_gvas before batch_copy (cherry picked from commit a80bd1c of fix-gva-invalid-v0.23, PR vllm-project#12643) Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
The store is already initialized eagerly during MemcacheBackend construction in _init_backend. Calling init_store() again in register_kv_caches is redundant; on the lazy path it re-triggers _setup_store()/SmemBmCreate2, which synchronously allocates the HBM/DRAM memory pool at startup and causes service-launch hangs on old HDK where memory allocation blocks. (cherry picked from commit 139a166 of fix-gva-invalid-v0.23, PR vllm-project#12643) Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
- _alloc_gvas_for_save: only cache and track gva>0 keys, skip failed batch_alloc - _prepare_load_gvas: only call batch_add_lease for keys with valid size (>0) - _prepare_load_gvas: report size<=0 and lease failure blocks to _invalid_block_ids - _prepare_load_gvas: skip invalid_block_ids for multi-group (hybrid) models - build_shared: filter out gva<=0 blocks from block_ids/block_gvas before batch_copy (cherry picked from commit a80bd1c of fix-gva-invalid-v0.23, PR vllm-project#12643) Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
The store is already initialized eagerly during MemcacheBackend construction in _init_backend. Calling init_store() again in register_kv_caches is redundant; on the lazy path it re-triggers _setup_store()/SmemBmCreate2, which synchronously allocates the HBM/DRAM memory pool at startup and causes service-launch hangs on old HDK where memory allocation blocks. (cherry picked from commit 139a166 of fix-gva-invalid-v0.23, PR vllm-project#12643) Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
- _alloc_gvas_for_save: only cache and track gva>0 keys, skip failed batch_alloc - _prepare_load_gvas: only call batch_add_lease for keys with valid size (>0) - _prepare_load_gvas: report size<=0 and lease failure blocks to _invalid_block_ids - _prepare_load_gvas: skip invalid_block_ids for multi-group (hybrid) models - build_shared: filter out gva<=0 blocks from block_ids/block_gvas before batch_copy (cherry picked from commit a80bd1c of fix-gva-invalid-v0.23, PR vllm-project#12643) Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
The store is already initialized eagerly during MemcacheBackend construction in _init_backend. Calling init_store() again in register_kv_caches is redundant; on the lazy path it re-triggers _setup_store()/SmemBmCreate2, which synchronously allocates the HBM/DRAM memory pool at startup and causes service-launch hangs on old HDK where memory allocation blocks. (cherry picked from commit 139a166 of fix-gva-invalid-v0.23, PR vllm-project#12643) Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
…parallel `_extract_physical_layer_index` returns a GLOBAL physical layer index (parsed from the model layer name), but the bound `>= self.num_layers` compares against the per-PP-stage layer count. With PP>1, stages beyond the first skip ALL their layers (global index >= local count), producing empty group metadata and crashing register_kv_caches/_init_layerwise_config. Use the model's total `num_hidden_layers` as the bound so the global index is correctly bounded. This only changes behavior when PP>1 (no effect when PP=1, since num_layers == num_hidden_layers; MTP is handled via speculative config and is not part of total_num_hidden_layers). Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
The two getattr() calls added for the PP layer-bound fix were missing the space after each comma, which `ruff format` (v0.14.0, pre-commit --all-files in CI) flags. Add the required spaces. Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
- _alloc_gvas_for_save: only cache and track gva>0 keys, skip failed batch_alloc - _prepare_load_gvas: only call batch_add_lease for keys with valid size (>0) - _prepare_load_gvas: report size<=0 and lease failure blocks to _invalid_block_ids - _prepare_load_gvas: skip invalid_block_ids for multi-group (hybrid) models - build_shared: filter out gva<=0 blocks from block_ids/block_gvas before batch_copy (cherry picked from commit a80bd1c of fix-gva-invalid-v0.23, PR vllm-project#12643) Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
The store is already initialized eagerly during MemcacheBackend construction in _init_backend. Calling init_store() again in register_kv_caches is redundant; on the lazy path it re-triggers _setup_store()/SmemBmCreate2, which synchronously allocates the HBM/DRAM memory pool at startup and causes service-launch hangs on old HDK where memory allocation blocks. (cherry picked from commit 139a166 of fix-gva-invalid-v0.23, PR vllm-project#12643) Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
- _alloc_gvas_for_save: only cache and track gva>0 keys, skip failed batch_alloc - _prepare_load_gvas: only call batch_add_lease for keys with valid size (>0) - _prepare_load_gvas: report size<=0 and lease failure blocks to _invalid_block_ids - _prepare_load_gvas: skip invalid_block_ids for multi-group (hybrid) models - build_shared: filter out gva<=0 blocks from block_ids/block_gvas before batch_copy (cherry picked from commit a80bd1c of fix-gva-invalid-v0.23, PR vllm-project#12643) Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
The store is already initialized eagerly during MemcacheBackend construction in _init_backend. Calling init_store() again in register_kv_caches is redundant; on the lazy path it re-triggers _setup_store()/SmemBmCreate2, which synchronously allocates the HBM/DRAM memory pool at startup and causes service-launch hangs on old HDK where memory allocation blocks. (cherry picked from commit 139a166 of fix-gva-invalid-v0.23, PR vllm-project#12643) Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
Cherry-pick of PR #12643 (fix-gva-invalid-v0.23) onto `main`. ## Changes - kv_transfer.py: _build_transfer_arrays warns on zero/negative base_gvas; build_shared filters out gva<=0 blocks before batch_copy. - pool_worker.py: _alloc_gvas_for_save only caches/tracks gva>0 keys; _prepare_load_gvas calls batch_add_lease only for valid keys, reports invalid/lease-failed blocks to _invalid_block_ids (single-group only), early-returns when block_ids_by_group is None. - pool_worker.py: remove redundant init_store() call in register_kv_caches (store already initialized eagerly in _init_backend; avoids re-triggering SmemBmCreate2 memory allocation at startup). ## Adaptation On `main` the _alloc_gvas_for_save path no longer collects all_group_save_keys (hybrid GDN refactor), so the save-path hunk was ported keeping only the `if gva > 0:` guard; the other 5 hunks apply verbatim. (cherry picked from commits a80bd1c / 139a166 of fix-gva-invalid-v0.23) - vLLM version: v0.25.1 - vLLM main: vllm-project/vllm@54503ec --------- Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
…2703) Cherry-pick of PR #12643 (fix-gva-invalid-v0.23) onto releases/v0.24.0rc. ## Changes - kv_transfer.py: _build_transfer_arrays warns on zero/negative base_gvas; build_shared filters out gva<=0 blocks before batch_copy. - pool_worker.py: _alloc_gvas_for_save only caches/tracks gva>0 keys; _prepare_load_gvas calls batch_add_lease only for valid keys, reports invalid/lease-failed blocks to _invalid_block_ids (single-group only), early-returns when block_ids_by_group is None. - pool_worker.py: remove redundant init_store() call in register_kv_caches (store already initialized eagerly in _init_backend; avoids re-triggering SmemBmCreate2 memory allocation at startup). ## Adaptation On releases/v0.24.0rc the _alloc_gvas_for_save path no longer collects all_group_save_keys (refactored by hybrid GDN feature), so the save-path hunk was ported keeping only the `if gva > 0:` guard; the other 5 hunks apply verbatim. (cherry picked from commits a80bd1c / 139a166 of fix-gva-invalid-v0.23) - vLLM version: v0.24.0 - vLLM main: vllm-project/vllm@85c09e9 --------- Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
…oject#12705) Cherry-pick of PR vllm-project#12643 (fix-gva-invalid-v0.23) onto `main`. ## Changes - kv_transfer.py: _build_transfer_arrays warns on zero/negative base_gvas; build_shared filters out gva<=0 blocks before batch_copy. - pool_worker.py: _alloc_gvas_for_save only caches/tracks gva>0 keys; _prepare_load_gvas calls batch_add_lease only for valid keys, reports invalid/lease-failed blocks to _invalid_block_ids (single-group only), early-returns when block_ids_by_group is None. - pool_worker.py: remove redundant init_store() call in register_kv_caches (store already initialized eagerly in _init_backend; avoids re-triggering SmemBmCreate2 memory allocation at startup). ## Adaptation On `main` the _alloc_gvas_for_save path no longer collects all_group_save_keys (hybrid GDN refactor), so the save-path hunk was ported keeping only the `if gva > 0:` guard; the other 5 hunks apply verbatim. (cherry picked from commits a80bd1c / 139a166 of fix-gva-invalid-v0.23) - vLLM version: v0.25.1 - vLLM main: vllm-project/vllm@54503ec --------- Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
…oject#12705) Cherry-pick of PR vllm-project#12643 (fix-gva-invalid-v0.23) onto `main`. ## Changes - kv_transfer.py: _build_transfer_arrays warns on zero/negative base_gvas; build_shared filters out gva<=0 blocks before batch_copy. - pool_worker.py: _alloc_gvas_for_save only caches/tracks gva>0 keys; _prepare_load_gvas calls batch_add_lease only for valid keys, reports invalid/lease-failed blocks to _invalid_block_ids (single-group only), early-returns when block_ids_by_group is None. - pool_worker.py: remove redundant init_store() call in register_kv_caches (store already initialized eagerly in _init_backend; avoids re-triggering SmemBmCreate2 memory allocation at startup). ## Adaptation On `main` the _alloc_gvas_for_save path no longer collects all_group_save_keys (hybrid GDN refactor), so the save-path hunk was ported keeping only the `if gva > 0:` guard; the other 5 hunks apply verbatim. (cherry picked from commits a80bd1c / 139a166 of fix-gva-invalid-v0.23) - vLLM version: v0.25.1 - vLLM main: vllm-project/vllm@54503ec --------- Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
…oject#12705) Cherry-pick of PR vllm-project#12643 (fix-gva-invalid-v0.23) onto `main`. ## Changes - kv_transfer.py: _build_transfer_arrays warns on zero/negative base_gvas; build_shared filters out gva<=0 blocks before batch_copy. - pool_worker.py: _alloc_gvas_for_save only caches/tracks gva>0 keys; _prepare_load_gvas calls batch_add_lease only for valid keys, reports invalid/lease-failed blocks to _invalid_block_ids (single-group only), early-returns when block_ids_by_group is None. - pool_worker.py: remove redundant init_store() call in register_kv_caches (store already initialized eagerly in _init_backend; avoids re-triggering SmemBmCreate2 memory allocation at startup). ## Adaptation On `main` the _alloc_gvas_for_save path no longer collects all_group_save_keys (hybrid GDN refactor), so the save-path hunk was ported keeping only the `if gva > 0:` guard; the other 5 hunks apply verbatim. (cherry picked from commits a80bd1c / 139a166 of fix-gva-invalid-v0.23) - vLLM version: v0.25.1 - vLLM main: vllm-project/vllm@54503ec --------- Signed-off-by: tyy0829 <87685049+tyy0829@users.noreply.github.com>
What this PR does
Fixes invalid GVA (value <= 0) handling when
batch_alloc/batch_get_key_info/batch_add_leasefail in the layerwise KV pool path.Changes
pool_worker.py
_alloc_gvas_for_save: Only cache and trackgva > 0keys; skip failedbatch_allocresults_prepare_load_gvas: Only callbatch_add_leasefor keys with validsize > 0_prepare_load_gvas: Reportsize <= 0and lease failure blocks to_invalid_block_ids_prepare_load_gvas: Skip_invalid_block_idsfor multi-group (hybrid) models to avoid scheduler crashkv_transfer.py
build_shared: Filter outgva <= 0blocks fromblock_ids_arr/block_gvas_arrbeforebatch_copy_build_transfer_arrays: Add warning check for zero/negativebase_gvasRationale
When
batch_allocorbatch_get_key_inforeturns GVA=0, the value flows through tobatch_copyas a size-like number (offset only), causing-3101errors. This PR ensures:batch_add_leaseis only called for valid keysbatch_copynever receives invalid GVA entriesReference: PR #9701 for
_invalid_block_idsmechanismSigned-off-by: tyy0829 tyy0829@users.noreply.github.com