From ec4560304b79d856acc9410f9d23f25cfeb0fe73 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhang <2448370773@qq.com> Date: Thu, 7 May 2026 15:37:53 +0200 Subject: [PATCH] [Bug Fix] Preserve decode state across retract-resume of GLM-5.1 (#23346) Co-authored-by: Shangming Cai --- python/sglang/srt/disaggregation/utils.py | 20 ++++----- python/sglang/srt/mem_cache/memory_pool.py | 48 ++++++++++++++++++++++ 2 files changed, 58 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 0bd5b5b76..26817b3de 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -243,16 +243,16 @@ class MetadataBuffers: def get_buf(self, idx: int): return ( - self.output_ids[idx], - self.cached_tokens[idx], - self.output_token_logprobs_val[idx], - self.output_token_logprobs_idx[idx], - self.output_top_logprobs_val[idx], - self.output_top_logprobs_idx[idx], - self.output_topk_p[idx], - self.output_topk_index[idx], - self.output_hidden_states[idx], - self.bootstrap_room[idx], + self.output_ids[idx].clone(), + self.cached_tokens[idx].clone(), + self.output_token_logprobs_val[idx].clone(), + self.output_token_logprobs_idx[idx].clone(), + self.output_top_logprobs_val[idx].clone(), + self.output_top_logprobs_idx[idx].clone(), + self.output_topk_p[idx].clone(), + self.output_topk_index[idx].clone(), + self.output_hidden_states[idx].clone(), + self.bootstrap_room[idx].clone(), ) def set_buf(self, req: Req): diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 6cbf67e3b..5d62bacb0 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1937,6 +1937,10 @@ class NSATokenToKVPool(MLATokenToKVPool): ] self._finalize_allocation_log(size) + def _clear_buffers(self): + del self.kv_buffer + del self.index_k_with_scale_buffer + def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor: if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) @@ -2011,6 +2015,50 @@ class NSATokenToKVPool(MLATokenToKVPool): pool=self, buf=buf, loc=loc, index_k=index_k, index_k_scale=index_k_scale ) + 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} + + def load_cpu_copy(self, kv_cache_cpu_dict, indices): + super().load_cpu_copy(kv_cache_cpu_dict["kv"], indices) + + page_indices = indices[:: self.page_size] // self.page_size + index_k_cpu = kv_cache_cpu_dict["index_k"] + torch.cuda.synchronize() + 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): + for i in range(0, len(page_indices), page_chunk_size): + chunk_page_indices = page_indices[i : i + page_chunk_size] + idx_cpu = index_k_cpu[layer_id][i // page_chunk_size] + assert idx_cpu.shape[0] == len(chunk_page_indices) + idx_chunk = idx_cpu.to( + self.index_k_with_scale_buffer[0].device, non_blocking=True + ) + self.index_k_with_scale_buffer[layer_id][chunk_page_indices] = idx_chunk + torch.cuda.synchronize() + def get_state_buf_infos(self): data_ptrs = [ self.index_k_with_scale_buffer[i].data_ptr() for i in range(self.layer_num)