Build the unified read stream directly, without the page-table rectangle (#37512)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

test/registered/unit/mem_cache/test_kv_index_translator.py

@@ -48,7 +48,7 @@ from types import SimpleNamespace
 import torch
 from test_multi_ended_allocator import _FakeUnifiedSWAKVPool
 
-from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
+from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator, KVReadTables
 from sglang.srt.mem_cache.multi_ended_allocator import (
     UnifiedSWATokenToKVPoolAllocator,
 )
@@ -235,6 +235,40 @@ class TestReadTableBuild(unittest.TestCase):
                     f"swa read table off-formula (ps={ps}, mult={swa_mult})",
                 )
 
+    def test_packed_stream_equals_the_rectangle_it_replaces(self):
+        """The packed builder and the rectangle must agree element for element:
+        packed[indptr[b] + p] == ids[b, p // ps] * ps + p % ps. Consumers that
+        plan over the stream and consumers that read the page table have to see
+        the same KV, so a change to either builder alone turns this red."""
+        for ps in (1, 4):
+            allocator = _build_composite(ps)
+            req_to_token, rows, seq_lens = _alloc_and_fill(
+                allocator, ps, lens=[5 * ps, 2 * ps, 3 * ps - 1]
+            )
+            src = _make_source(allocator, req_to_token, ps)
+            view = src.build_index_table(
+                req_pool_indices=rows, seq_lens=seq_lens, max_pages=6
+            )
+            indptr = torch.zeros(len(seq_lens) + 1, dtype=torch.int32)
+            indptr[1:] = torch.cumsum(seq_lens, dim=0)
+            total = int(indptr[-1])
+            packed = torch.zeros(total, dtype=torch.int32)
+            translated = src.fill_packed_read_stream(
+                req_pool_indices=rows,
+                seq_lens=seq_lens,
+                indptr=indptr,
+                total_tokens=total,
+                out=packed,
+            )
+            self.assertTrue(translated)
+            for b, n in enumerate(seq_lens.tolist()):
+                for pos in range(n):
+                    self.assertEqual(
+                        int(packed[int(indptr[b]) + pos]),
+                        int(view.ids[b, pos // ps]) * ps + pos % ps,
+                        f"packed stream off the page table (ps={ps}, b={b}, {pos=})",
+                    )
+
     def test_sink_routing(self):
         """Dead lanes (seq_len 0), -1 slots inside the live prefix, and
         tombstoned v2p pages must ALL read entry 0 -- one wild entry is a
@@ -426,10 +460,11 @@ class TestCaptureContract(unittest.TestCase):
         req_to_token[1, : 2 * ps] = v
         src = _make_source(allocator, req_to_token, ps)
 
-        tables = src.make_capture_tables(max_bs=4, max_context_len=8 * ps)
-        cap, cap_swa = tables.full, tables.sliding_window
-        self.assertTrue(bool((cap == 0).all()), "read tables must start zeroed")
-        self.assertIsNotNone(cap_swa, "the SWA composite has a second id space")
+        # A page-table consumer owns its buffers; entry 0 is the reserved sink
+        # in every id space, so zeros are what an unfilled column must read as.
+        cap = torch.zeros((4, 8), dtype=torch.int32, device=_DEV)
+        cap_swa = torch.zeros((4, 8), dtype=torch.int32, device=_DEV)
+        tables = KVReadTables(full=cap, sliding_window=cap_swa)
 
         # Poison everything, then refresh a 1-row batch: ONLY its live prefix
         # may change -- stale tails and other rows are the fa3 contract.

test/registered/unit/mem_cache/test_unified_mla_block_table.py

@@ -261,32 +261,32 @@ class TestFa3MetadataBlockTable(unittest.TestCase):
         max_seq_pages = (int(sl.max().item()) + page_size - 1) // page_size
 
         if v2p:
-            # The translator's canonical, then the wrapper copies its rows.
-            canonical = torch.zeros((bs, max_pages), dtype=torch.int32, device=_DEV)
-            build_kv_read_table(
-                req_to_token=rt,
-                req_pool_indices=rpi,
-                seq_lens=sl.to(torch.int64),
-                v2p=v2p_full,
-                multiplier=mult,
-                page_size=page_size,
-                max_pages=max_pages,
-                out=canonical,
-            )
-            rows = torch.arange(bs, dtype=torch.int64, device=_DEV)
+            # Unified: the fused call does the prefix sum only, and the builder
+            # writes the page table itself -- bounded by the `cache_seqlens`
+            # that call just produced, which is what a reader bounds by too.
             normal_decode_set_metadata(
                 cache_seqlens,
                 cu_seqlens_k,
                 page_table,
-                canonical,
-                rows,
+                rt,
+                rpi,
                 max_seq_pages,
                 sl.to(torch.int64),
                 0,
                 page_size,
                 None,
                 None,
-                src_is_read_table=True,
+                skip_page_table=True,
+            )
+            build_kv_read_table(
+                req_to_token=rt,
+                req_pool_indices=rpi,
+                seq_lens=cache_seqlens,
+                v2p=v2p_full,
+                multiplier=mult,
+                page_size=page_size,
+                max_pages=max_pages,
+                out=page_table,
             )
         else:
             normal_decode_set_metadata(