fix(req_pool): bump pool.size to match actual tensor row count after #24243 (#24439)

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()
 
 

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:

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,

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,
             ),