[Fix] Write the req_to_token page tail so rows stay valid over whole pages

python/sglang/srt/disaggregation/decode.py

@@ -70,6 +70,7 @@
 )
 from sglang.srt.managers.schedule_policy import match_prefix_for_req
 from sglang.srt.managers.utils import GenerationBatchResult
+from sglang.srt.mem_cache.allocation import write_page_tail
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams
 from sglang.srt.mem_cache.common import (
@@ -1709,13 +1710,18 @@ def _pre_alloc(
             f"req={req.rid}"
         )
 
+        write_end = total_prefix_len + len(kv_loc)
         self.req_to_token_pool.write(
-            (
-                req.req_pool_idx,
-                slice(total_prefix_len, total_prefix_len + len(kv_loc)),
-            ),
+            (req.req_pool_idx, slice(total_prefix_len, write_end)),
             kv_loc,
         )
+        write_page_tail(
+            allocator,
+            self.req_to_token_pool.req_to_token,
+            torch.tensor([req.req_pool_idx], device=kv_loc.device),
+            torch.tensor([write_end], device=kv_loc.device),
+            allocator.page_size,
+        )
 
         # Truncate fill_len to kv_committed_len so cache_unfinished_req only
         # inserts committed KV into the radix tree. The last output token

python/sglang/srt/mem_cache/allocation.py

@@ -103,6 +103,133 @@ def write_cache_indices(
             pt += extend_len
 
 
+def write_page_tail_indices(
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    write_ends: torch.Tensor,
+    page_size: int,
+) -> None:
+    # The row write stops at the requested tokens, so the tail of the last page
+    # -- allocated, and consecutive with the last slot -- keeps stale indices.
+    if page_size == 1:
+        return
+
+    device = req_to_token.device
+    steps = torch.arange(page_size - 1, device=device, dtype=write_ends.dtype)
+    last_pos = (write_ends - 1).clamp(min=0)
+    last_val = req_to_token[req_pool_indices, last_pos]
+
+    positions = write_ends[:, None] + steps[None, :]
+    ceilings = ((write_ends + page_size - 1) // page_size * page_size)[:, None]
+    in_page = positions < ceilings
+
+    # Lanes past the ceiling fold onto the last real entry and rewrite it with
+    # its own value, keeping the scatter a fixed (bs, page_size - 1) shape.
+    positions = torch.where(in_page, positions, last_pos[:, None])
+    values = torch.where(
+        in_page,
+        last_val[:, None] + steps[None, :].to(last_val.dtype) + 1,
+        last_val[:, None],
+    )
+    req_to_token[req_pool_indices[:, None].expand_as(positions), positions] = values
+
+
+def write_page_tail_mapping(
+    allocator,
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    write_ends: torch.Tensor,
+    page_size: int,
+) -> None:
+    """The SWA counterpart of write_page_tail_indices: publish full->swa for the
+    rest of the last page, so a decode step landing there finds a live mapping."""
+    if page_size == 1:
+        return
+
+    device = req_to_token.device
+    steps = torch.arange(page_size - 1, device=device, dtype=write_ends.dtype)
+    last_pos = (write_ends - 1).clamp(min=0)
+    last_full = req_to_token[req_pool_indices, last_pos].to(torch.int64)
+    last_swa = allocator.full_to_swa_index_mapping[last_full].to(torch.int64)
+
+    ceilings = ((write_ends + page_size - 1) // page_size * page_size)[:, None]
+    in_page = write_ends[:, None] + steps[None, :] < ceilings
+    # Lanes past the ceiling re-publish the last real pair, an idempotent write.
+    offsets = steps[None, :] + 1
+    full_indices = torch.where(
+        in_page, last_full[:, None] + offsets, last_full[:, None]
+    )
+    swa_indices = torch.where(in_page, last_swa[:, None] + offsets, last_swa[:, None])
+    allocator.set_full_to_swa_mapping(full_indices.reshape(-1), swa_indices.reshape(-1))
+
+
+def write_page_tail(
+    allocator,
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    write_ends: torch.Tensor,
+    page_size: int,
+) -> None:
+    """Publish the rest of the last page: row indices always, plus the full->swa
+    mapping when the allocator keeps one."""
+    from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
+
+    write_page_tail_indices(req_to_token, req_pool_indices, write_ends, page_size)
+    if isinstance(allocator, SWATokenToKVPoolAllocator):
+        write_page_tail_mapping(
+            allocator, req_to_token, req_pool_indices, write_ends, page_size
+        )
+
+
+def uses_page_granular_decode(batch: ScheduleBatch, token_per_req: int) -> bool:
+    # Exact type, not isinstance: DSV4 bundles and hisparse device buffers do
+    # per-step work beyond handing out an index, so they keep the per-token path.
+    from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator
+    from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
+
+    return (
+        _alloc_page_size(batch) > 1
+        and token_per_req == 1
+        and not batch.model_config.is_encoder_decoder
+        and type(batch.tree_cache.token_to_kv_pool_allocator)
+        in (PagedTokenToKVPoolAllocator, SWATokenToKVPoolAllocator)
+    )
+
+
+def publish_decode_pages(batch: ScheduleBatch, page_size: int) -> None:
+    """Hand a whole page to every request that opens one this step and write it
+    into the request's row, so the steps inside that page allocate nothing."""
+    openers = [
+        i for i, n in enumerate(batch.seq_lens_cpu.tolist()) if n % page_size == 0
+    ]
+    if not openers:
+        return
+
+    allocator = batch.tree_cache.token_to_kv_pool_allocator
+    num_tokens = len(openers) * page_size
+    evict_from_tree_cache(batch.tree_cache, num_tokens)
+    pages = allocator.alloc(num_tokens)
+    if pages is None:
+        error_msg = (
+            f"Decode out of memory. Try to lower your batch size.\n"
+            f"Try to allocate {num_tokens} tokens.\n"
+            f"{available_and_evictable_str(batch.tree_cache)}"
+        )
+        logger.error(error_msg)
+        batch.tree_cache.pretty_print()
+        raise RuntimeError(error_msg)
+
+    device = batch.device
+    opener_indices = torch.tensor(openers, dtype=torch.int64, device=device)
+    rows = batch.req_pool_indices[opener_indices]
+    starts = batch.seq_lens[opener_indices]
+    positions = starts[:, None] + torch.arange(page_size, device=device)
+    batch.req_to_token_pool.write(
+        (rows[:, None].expand_as(positions), positions),
+        pages.view(len(openers), page_size).to(torch.int32),
+    )
+
+
 def get_last_loc(
     req_to_token: torch.Tensor,
     req_pool_indices_tensor: torch.Tensor,
@@ -359,6 +486,13 @@ def alloc_for_extend(
         prefix_tensors,
         batch.req_to_token_pool,
     )
+    write_page_tail(
+        batch.tree_cache.token_to_kv_pool_allocator,
+        batch.req_to_token_pool.req_to_token,
+        req_pool_indices_device,
+        batch.seq_lens,
+        alloc_page_size,
+    )
 
     # DSV4-NPU hook: no-op on non-DSV4 paths.
     if _is_npu:
@@ -522,7 +656,15 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
     seq_lens_gpu = batch.seq_lens
     bs = seq_lens_gpu.shape[0]
 
-    if _alloc_page_size(batch) == 1:
+    page_granular = uses_page_granular_decode(batch, token_per_req)
+    if page_granular:
+        # The page is published whole when it opens, so this step's slot is
+        # already in the row -- gather it instead of allocating one.
+        publish_decode_pages(batch, _alloc_page_size(batch))
+        out_cache_loc = batch.req_to_token_pool.req_to_token[
+            batch.req_pool_indices, seq_lens_gpu
+        ].to(torch.int64)
+    elif _alloc_page_size(batch) == 1:
         # Non-paged allocation
         out_cache_loc = alloc_token_slots(batch.tree_cache, bs * token_per_req)
     else:
@@ -541,24 +683,25 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
             batch=batch,
         )
 
-    # Write to req_to_token_pool
-    if batch.model_config.is_encoder_decoder:
-        locs = batch.encoder_lens + seq_lens_gpu
-    else:
-        locs = seq_lens_gpu.clone()
-
-    batch.req_to_token_pool.write(
-        (batch.req_pool_indices, locs), out_cache_loc.to(torch.int32)
-    )
+    if not page_granular:
+        # Write to req_to_token_pool
+        if batch.model_config.is_encoder_decoder:
+            locs = batch.encoder_lens + seq_lens_gpu
+        else:
+            locs = seq_lens_gpu.clone()
 
-    # DSV4-NPU hook: no-op on non-DSV4 paths.
-    if _is_npu:
-        maybe_write_dsv4_decode(
-            batch,
-            batch.seq_lens_cpu + token_per_req,
-            token_per_req,
+        batch.req_to_token_pool.write(
+            (batch.req_pool_indices, locs), out_cache_loc.to(torch.int32)
         )
 
+        # DSV4-NPU hook: no-op on non-DSV4 paths.
+        if _is_npu:
+            maybe_write_dsv4_decode(
+                batch,
+                batch.seq_lens_cpu + token_per_req,
+                token_per_req,
+            )
+
     for req in batch.reqs:
         req.kv.kv_allocated_len += token_per_req
 

python/sglang/srt/mem_cache/allocation_sizing.py

@@ -80,17 +80,18 @@ def page_aligned_decode_alloc_lens(
 def get_req_to_token_extra_context_len() -> int:
     """req_to_token row headroom beyond the model context length.
 
-    Sized to hold the decode over-allocation; the spec v2 page>1 topk>1 holey
-    draft footprint can outgrow the default num_draft_tokens headroom. The row
-    headroom and the pools it sits next to derive from the same bag leaves, so
-    they cannot disagree after a post-publish override.
+    Sized to hold the page-ceiled row write and the decode over-allocation; the
+    spec v2 page>1 topk>1 holey draft footprint can outgrow the default
+    num_draft_tokens headroom. The row headroom and the pools it sits next to
+    derive from the same bag leaves, so they cannot disagree after a
+    post-publish override.
     """
-    # FIXME(lsyin): temporary fix for the context length issue under spec decoding
     extra = 4 + (max_speculative_num_draft_tokens() or 0)
     page_size = get_alloc_page_size()
+    if page_size > 1:
+        # A request opening a page at the context limit publishes a whole page
+        # past it; without the headroom that write lands in the neighbor row.
+        extra = max(extra, page_size)
     if get_spec().speculative_algorithm is not None and page_size > 1:
-        # kv_allocated_len is page-aligned (eagle_prepare_for_decode), so near
-        # the context limit the aligned reserve can overshoot by page_size - 1;
-        # without the headroom the row write silently lands in the neighbor row.
         extra = max(extra, get_alloc_reserve_per_decode() + page_size - 1)
     return extra

python/sglang/srt/mem_cache/allocator/swa.py

@@ -150,7 +150,9 @@ def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
         return self._kvcache.translate_loc_from_full_to_swa(kv_indices)
 
     def alloc(self, need_size: int):
-        assert self.page_size == 1
+        # Whole pages at page_size > 1: both sub-allocators lay a page out in row
+        # order, so the two returns line up index-for-index and the mapping below
+        # publishes the whole page at once.
         if need_size > self.full_attn_allocator.available_size():
             return None
         if need_size > self.swa_attn_allocator.available_size():

test/registered/unit/mem_cache/test_hisparse_allocator.py

@@ -102,6 +102,7 @@ def set_extend_range(start, end):
         class ReqToTokenPool:
             def __init__(self):
                 self.writes = []
+                self.req_to_token = torch.zeros((1, 1024), dtype=torch.int32)
 
             def alloc(self, reqs):
                 for item in reqs:
@@ -110,6 +111,7 @@ def alloc(self, reqs):
 
             def write(self, indices, values):
                 self.writes.append((indices, values))
+                self.req_to_token[indices] = values.to(torch.int32)
 
         req_to_token_pool = ReqToTokenPool()
         allocator = SimpleNamespace(

test/registered/unit/mem_cache/test_page_tail_indices.py

@@ -0,0 +1,86 @@
+"""Unit tests for write_page_tail_indices: the req_to_token row must stay valid
+over every page the allocator handed out, not just over the requested tokens."""
+
+import unittest
+
+import torch
+
+from sglang.srt.mem_cache.allocation import write_page_tail_indices
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+_ROW_WIDTH = 64
+
+
+def _paged_row(page_ids, length, page_size, dtype):
+    # A row as the allocator lays it out: page_ids[k] backs row positions
+    # [k*ps, (k+1)*ps), and a slot sits at its own offset inside that page.
+    positions = torch.arange(length)
+    slots = page_ids[positions // page_size] * page_size + positions % page_size
+    return slots.to(dtype)
+
+
+class TestWritePageTailIndices(CustomTestCase):
+    def test_tail_continues_the_last_page(self):
+        rtt = torch.zeros((1, _ROW_WIDTH), dtype=torch.int32)
+        rtt[0, :5] = torch.arange(40, 45, dtype=torch.int32)
+
+        write_page_tail_indices(rtt, torch.tensor([0]), torch.tensor([5]), 4)
+
+        self.assertEqual(rtt[0, :8].tolist(), [40, 41, 42, 43, 44, 45, 46, 47])
+
+    def test_completes_the_last_page_and_moves_nothing_else(self):
+        # The oracle is the allocator's layout -- a slot sits at its own page
+        # offset, and the tail shares the page of the last written slot -- not
+        # the implementation's own "last index plus one" rule. Checking every
+        # row also subsumes empty rows, page-aligned ends and page_size 1.
+        generator = torch.Generator().manual_seed(0)
+        num_rows = 4
+        for page_size in (1, 2, 4, 8, 16):
+            num_pages = _ROW_WIDTH // page_size
+            for _ in range(50):
+                batch_size = int(
+                    torch.randint(1, num_rows + 1, (1,), generator=generator)
+                )
+                req_pool_indices = torch.randperm(num_rows, generator=generator)[
+                    :batch_size
+                ]
+                write_ends = torch.randint(
+                    0, _ROW_WIDTH - page_size + 1, (batch_size,), generator=generator
+                )
+                rtt = torch.randint(
+                    0, 1000, (num_rows, _ROW_WIDTH), generator=generator
+                ).to(torch.int32)
+                for row, end in zip(req_pool_indices.tolist(), write_ends.tolist()):
+                    page_ids = torch.randperm(64, generator=generator)[:num_pages]
+                    rtt[row, :end] = _paged_row(page_ids, end, page_size, rtt.dtype)
+                before = rtt.clone()
+
+                write_page_tail_indices(rtt, req_pool_indices, write_ends, page_size)
+
+                offsets = torch.arange(_ROW_WIDTH, dtype=rtt.dtype) % page_size
+                touched = req_pool_indices.tolist()
+                for row, end in zip(touched, write_ends.tolist()):
+                    ceiling = -(-end // page_size) * page_size
+                    where = f"{page_size=} {row=} {end=}"
+                    after, prior = rtt[row], before[row]
+
+                    page_offsets = after[:ceiling] % page_size
+                    tail_pages = after[end:ceiling] // page_size
+                    last_page = after[end - 1] // page_size
+
+                    self.assertTrue(torch.equal(page_offsets, offsets[:ceiling]), where)
+                    self.assertTrue(torch.all(tail_pages == last_page), where)
+                    self.assertTrue(torch.equal(after[:end], prior[:end]), where)
+                    self.assertTrue(
+                        torch.equal(after[ceiling:], prior[ceiling:]), where
+                    )
+
+                untouched = [r for r in range(num_rows) if r not in touched]
+                self.assertTrue(torch.equal(rtt[untouched], before[untouched]))
+
+
+if __name__ == "__main__":
+    unittest.main()