Repository navigation
[Bug Fix] Preserve decode state across retract-resume of GLM-5.1 - #23346
Conversation
![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 introduces several improvements to disaggregated execution, including correct metadata population for fill_ids, the addition of .clone() when retrieving buffers to ensure data integrity, and the implementation of CPU offloading for NSA-specific buffers. A review comment identifies an inconsistency in the get_cpu_copy return type within NSATokenToKVPool compared to other pool implementations, which could lead to issues if the system expects a standard list-of-lists format.
| def get_cpu_copy(self, indices): | ||
| # NSA keeps a page-indexed index_k_with_scale_buffer alongside kv_buffer. | ||
| # Retract frees the slots/pages and they get reused by other reqs' | ||
| # set_index_k_scale_buffer, so we must offload it here too -- otherwise | ||
| # resume restores kv_buffer but leaves foreign index/scale in place and | ||
| # NSA attention reads garbage at those token positions. | ||
| kv_cache_cpu = super().get_cpu_copy(indices) | ||
|
|
||
| page_indices = indices[:: self.page_size] // self.page_size | ||
| torch.cuda.synchronize() | ||
| index_k_cpu = [] | ||
| chunk_size = self.cpu_offloading_chunk_size | ||
| page_chunk_size = max(1, chunk_size // self.page_size) | ||
| for layer_id in range(self.layer_num): | ||
| index_k_cpu.append([]) | ||
| for i in range(0, len(page_indices), page_chunk_size): | ||
| chunk_page_indices = page_indices[i : i + page_chunk_size] | ||
| idx_cpu = self.index_k_with_scale_buffer[layer_id][ | ||
| chunk_page_indices | ||
| ].to("cpu", non_blocking=True) | ||
| index_k_cpu[-1].append(idx_cpu) | ||
| torch.cuda.synchronize() | ||
|
|
||
| return {"kv": kv_cache_cpu, "index_k": index_k_cpu} |
There was a problem hiding this comment.
The return type of get_cpu_copy in NSATokenToKVPool is changed to a dictionary, whereas the base class KVCache and sibling implementations (like MHATokenToKVPool and MLATokenToKVPool) return a List[List[Tensor]]. While load_cpu_copy is updated to handle this dictionary, this inconsistency might cause issues if other parts of the system expect the standard list-of-lists format for CPU-offloaded KV cache data. Consider if a more unified multi-buffer offloading interface should be introduced in the base class.
|
/tag-and-rerun-ci |
|
/rerun-stage stage-c-test-8-gpu-h20 |
|
✅ Triggered |
|
/rerun-stage stage-c-test-8-gpu-h20 |
|
✅ Triggered |
|
/tag-and-rerun-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
2 similar comments
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
The fill_len logic has been fixed by the decode radix cache PR: f589547 Due to the conflict, I have fixed it to adopt the incoming change from main, which is more robust. |
|
/rerun-stage stage-c-test-8-gpu-h20 |
|
✅ Triggered |
|
/rerun-failed-ci |
|
/rerun-test test/registered/disaggregation/test_disaggregation_basic.py |
|
✅ |
…-project#23346) Co-authored-by: Shangming Cai <csmthu@gmail.com>
…-project#23346) Co-authored-by: Shangming Cai <csmthu@gmail.com>
5 participants
PD Disaggregated Decode: Fix State Aliasing on Retract-Resume
GLM-5.1 on PD-disaggregated decode hits three related state-aliasing bugs when a retracted req resumes. Same pattern each time: a slot is freed and reused, but the original owner still references it. Resume reads foreign data and emits garbled output, which can further leak via the radix tree.
MetadataBuffers.get_bufReturn
.clone()instead of slice views. The shared slot gets overwritten by the nextset_buf; retract-resumed reqs runprocess_prebuiltmuch later and end up seeding EAGLE with foreign hidden states and topk.DecodePreallocQueue._pre_allocTrim
fill_idsso its length matchesfill_len. The last sampled token's KV is never offloaded but gets re-fed on resume, causing a read past the validreq_to_tokenrange. The stale index leaked into the radix tree and corrupted prefix-matched reqs. Added an assert for the invariant.GLM-5.1 carries an extra
index_k_with_scale_bufferalongsidekv_buffer. Overrideget_cpu_copy/load_cpu_copyto offload and restore both. Previously only KV was saved, so resume left foreign index/scale in place and NSA sparse indexing picked wrong keys. Also override_clear_buffersto release the extra buffer on teardown, matching sibling KVCache subclasses.