diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index ae666e35a..d2895242a 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -127,18 +127,19 @@ class DecodeReqToTokenPool: ) self.size = size + # +1 padding row at index 0; see ReqToTokenPool for rationale. + self._alloc_size = size + pre_alloc_size + 1 self.max_context_len = max_context_len self.device = device self.pre_alloc_size = pre_alloc_size with memory_saver_adapter.region(tag=GPU_MEMORY_TYPE_KV_CACHE): - # +1 row 0 padding; mirrors ReqToTokenPool / KV pool padding slot 0. self.req_to_token = torch.zeros( - (size + pre_alloc_size + 1, max_context_len), + (self._alloc_size, max_context_len), dtype=torch.int32, device=device, ) - self.free_slots = list(range(1, size + pre_alloc_size + 1)) + self.free_slots = list(range(1, self._alloc_size)) def write(self, indices, values): self.req_to_token[indices] = values @@ -175,7 +176,7 @@ class DecodeReqToTokenPool: req.req_pool_idx = None def clear(self): - self.free_slots = list(range(1, self.size + self.pre_alloc_size + 1)) + self.free_slots = list(range(1, self._alloc_size)) class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): @@ -240,7 +241,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): ) def clear(self): - self.free_slots = list(range(1, self.size + self.pre_alloc_size + 1)) + self.free_slots = list(range(1, self._alloc_size)) self.mamba_pool.clear() diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index a0e0765f2..6cbf67e3b 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -139,17 +139,16 @@ class ReqToTokenPool: ) self.size = size + # +1 padding row at index 0: cuda-graph padded batches default + # req_pool_indices to 0, so dummy reads/writes land here harmlessly. + self._alloc_size = size + 1 self.max_context_len = max_context_len self.device = device with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): - # +1 row for padding slot 0 (mirrors KV pool): cuda-graph padded - # batches default req_pool_indices to 0, so routing dummies through - # unowned slot 0 keeps req_to_token[0, :] zero and downstream writes - # harmless. self.req_to_token = torch.zeros( - (size + 1, max_context_len), dtype=torch.int32, device=device + (self._alloc_size, max_context_len), dtype=torch.int32, device=device ) - self.free_slots = list(range(1, size + 1)) + self.free_slots = list(range(1, self._alloc_size)) def write(self, indices, values): self.req_to_token[indices] = values @@ -189,7 +188,7 @@ class ReqToTokenPool: req.req_pool_idx = None def clear(self): - self.free_slots = list(range(1, self.size + 1)) + self.free_slots = list(range(1, self._alloc_size)) class MambaPool: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 820224ace..7bdd8ffd6 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2794,8 +2794,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): if self.use_ngram_embedding: from sglang.srt.layers.n_gram_embedding import NgramEmbedding + # Sized to mirror req_to_token (indexed by req_pool_idx). self.token_table = torch.empty( - self.req_to_token_pool.size, + self.req_to_token_pool.req_to_token.shape[0], self.model_config.context_len, dtype=torch.int32, device=self.device, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 9eac638ee..9372209b7 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -146,10 +146,11 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): self.init_lm_head() - # Used for KV Cache reversion + # KV cache reversion buffer; sized to mirror req_to_token (indexed by + # req_pool_idx). self.req_to_hidden_states_pool = torch.empty( ( - self.req_to_token_pool.size, + self.req_to_token_pool.req_to_token.shape[0], self.speculative_num_steps - 1, self.model_config.hidden_size, ),