[Feature] Unified memory: support decode context parallelism for the trtllm_mla family (#37693)

python/sglang/kernels/ops/attention/dcp_kernels.py

@@ -82,13 +82,23 @@ def create_mla_kv_page_table_for_dcp(
     req_pool_indices_ptr,
     local_seq_lens_ptr,
     block_kv_indices_ptr,
+    v2p_ptr,  # in: [num_pages + 1] int64 -- virtual->physical page table
     req_to_token_stride: tl.constexpr,
     block_table_stride: tl.constexpr,
+    mult,  # runtime: kernel_page_multiplier of the target sub-pool
     PHYSICAL_PAGE_SIZE: tl.constexpr,
     DCP_SIZE: tl.constexpr,
     DCP_RANK: tl.constexpr,
     PAGES_PER_BLOCK: tl.constexpr,
+    HAS_V2P: tl.constexpr,
 ):
+    """This rank's cyclic slice of each request, as a page table.
+
+    ``HAS_V2P`` picks the id space the emitted page number is in: the
+    DCP-collapsed page IS physical on a static pool, and still VIRTUAL under
+    the unified memory pool, where it takes one more gather through ``v2p_ptr``
+    and a ``mult`` scale to reach the per-layer views.
+    """
     req = tl.program_id(0)
     page_block = tl.program_id(1)
     page_offsets = page_block * PAGES_PER_BLOCK + tl.arange(0, PAGES_PER_BLOCK)
@@ -102,10 +112,16 @@ def create_mla_kv_page_table_for_dcp(
         mask=mask,
         other=0,
     )
-    physical_pages = virtual_locs // DCP_SIZE // PHYSICAL_PAGE_SIZE
+    pages = virtual_locs // DCP_SIZE // PHYSICAL_PAGE_SIZE
+    if HAS_V2P:
+        # A `-1` in req_to_token and a freed (`-1`) v2p row both clamp to entry
+        # 0, the reserved padding page.
+        pages = tl.where(virtual_locs < 0, 0, pages)
+        physical = tl.load(v2p_ptr + pages, mask=mask, other=0)
+        pages = tl.maximum(physical * mult, 0)
     tl.store(
         block_kv_indices_ptr + req * block_table_stride + page_offsets,
-        physical_pages,
+        pages.to(tl.int32),
         mask=mask,
     )
 

python/sglang/srt/arg_groups/kv_cache_hook.py

@@ -305,12 +305,10 @@ def _validate_unified_memory_dcp(server_args: Any) -> None:
         "transfer, where translate_kv_indices_for_transfer would abort a "
         "server that had already booted."
     )
-    # trtllm_mla (and its cutedsl_mla / tokenspeed_mla subclasses) build the
+    # The trtllm_mla family builds its DCP block table through the pool's v2p
-    # MLA block table straight from req_to_token with
+    # gather (create_mla_kv_page_table_for_dcp), so it speaks the same
-    # create_flashmla_kv_indices_triton, whose v2p gather assumes UNWIDENED
+    # two-stage contract as flashinfer.
-    # page ids; the DCP variant (create_mla_kv_page_table_for_dcp) has no v2p
+    dcp_allowed = {"flashinfer", "trtllm_mla", "cutedsl_mla", "tokenspeed_mla"}
-    # gather at all. Wire one of them through the other to add those here.
-    dcp_allowed = {"flashinfer"}
     backends = set(attention_backends_of(resolved_view(server_args)))
     backends.discard(None)
     assert backends <= dcp_allowed, (

python/sglang/srt/layers/attention/cutedsl_mla_backend.py

@@ -174,17 +174,12 @@ class CuteDslMLABackend(TRTLLMMLABackend):
         if self.data_type == torch.float8_e4m3fn:
             assert q_rope is not None and k_rope is not None
             if cos_sin_cache is None:
-                if (
+                if save_kv_cache and self._fused_set_kv_concat_q_fp8:
-                    save_kv_cache
+                    loc = self._resolve_fused_write_loc(forward_batch)
-                    and self._fused_set_kv_concat_q_fp8
+                    if loc is not None:
-                    and not self.kv_index_translator.is_translating
-                ):
-                    # Static pool: out_cache_loc is already the physical loc.
-                    # Fused: bf16->fp8 quantize + KV scatter + q concat in one
-                    # launch; None when not covered.
                         query = self._set_kv_and_concat_q_fp8_fused(
                             layer=layer,
-                        loc=forward_batch.out_cache_loc,
+                            loc=loc,
                             q=q,
                             q_rope=q_rope,
                             k=k,
@@ -211,7 +206,7 @@ class CuteDslMLABackend(TRTLLMMLABackend):
         if query is None and save_kv_cache:
             assert k is not None and k_rope is not None
             self.token_to_kv_pool.set_mla_kv_buffer(
-                layer, forward_batch.out_cache_loc, k, k_rope
+                layer, self._kv_write_loc(forward_batch), k, k_rope
             )
 
         if query is not None:

python/sglang/srt/layers/attention/tokenspeed_mla_backend.py

@@ -389,7 +389,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
 
         if save_kv_cache:
             self.token_to_kv_pool.set_mla_kv_buffer(
-                layer, forward_batch.out_cache_loc, k, k_rope
+                layer, self._kv_write_loc(forward_batch), k, k_rope
             )
 
         query = q.view(-1, layer.tp_q_head_num, layer.head_dim)

python/sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -377,6 +377,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
     ) -> None:
         parallel = get_parallel()
         pages_per_block = get_num_page_per_block_flashmla(self.page_size)
+        # None on a static pool, whose collapsed page is already physical.
+        v2p = self.kv_index_translator.full_v2p_table
         create_mla_kv_page_table_for_dcp[
             (
                 block_kv_indices.shape[0],
@@ -389,12 +391,15 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             req_pool_indices,
             local_seq_lens,
             block_kv_indices,
+            v2p,
             self.req_to_token.stride(0),
             block_kv_indices.stride(0),
+            self.kv_index_translator.full_page_multiplier,
             PHYSICAL_PAGE_SIZE=self.page_size,
             DCP_SIZE=parallel.dcp_size,
             DCP_RANK=parallel.dcp_rank,
             PAGES_PER_BLOCK=pages_per_block,
+            HAS_V2P=v2p is not None,
         )
 
     def _create_block_kv_indices(
@@ -777,6 +782,16 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         else:
             self._decode_kernel_loc = None
 
+    def _kv_write_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
+        """The loc an unfused KV scatter must write at: the capture-stable
+        buffer under a captured unified-pool decode, since the translate
+        rebinds `out_cache_loc` to a fresh tensor the graph never recorded;
+        the batch's own loc everywhere else.
+        """
+        if self._decode_kernel_loc is not None:
+            return self._decode_kernel_loc
+        return forward_batch.out_cache_loc
+
     def _resolve_fused_write_loc(
         self, forward_batch: ForwardBatch
     ) -> Optional[torch.Tensor]:
@@ -1198,6 +1213,12 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         ):
             return None
         parallel = get_parallel()
+        # `loc` is WIDENED: the kernel resolves the owner rule itself, and that
+        # is also its only skip. A DCP-resolved loc never reaches here -- see
+        # the `_fused_set_kv_concat_q_fp8` gate.
+        assert not (parallel.dcp_enabled and self.kv_index_translator.is_translating), (
+            "fused fp8 KV write reached with a DCP-resolved loc"
+        )
         return set_mla_kv_concat_q_fp8(
             kv_buffer=kv_2d,
             loc=loc,
@@ -1205,8 +1226,6 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             cache_k_rope=k_rope_2d,
             q_nope=q_nope,
             q_rope=q_rope_3d,
-            # DCP cyclic KV sharding: virtual loc -> owner mask + loc//world
-            # (identity when attn_dcp_size == 1).
             dcp_world_size=parallel.attn_dcp_size,
             dcp_rank=parallel.attn_dcp_rank,
         )

python/sglang/srt/mem_cache/kv_index_translator.py

@@ -392,6 +392,22 @@ class KVIndexTranslator:
         self._index_table_memo = (weakref.ref(forward_batch), view)
         return view
 
+    @property
+    def full_v2p_table(self) -> Optional[torch.Tensor]:
+        """The full-attention virtual->physical PAGE table, or None when this
+        pool needs no translation.
+
+        For the DCP page-table builders, whose gather is over a rank's cyclic
+        slice rather than a row prefix, so `build_index_table` cannot serve
+        them.
+        """
+        return self._full_v2p_table
+
+    @property
+    def full_page_multiplier(self) -> int:
+        """Scales a physical page into the id space the per-layer views use."""
+        return self._full_page_multiplier
+
     def bind_and_verify_backends(self, backends) -> None:
         """Boot: make every reachable backend carry THIS translator.
 

test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py

@@ -213,6 +213,225 @@ class TestTRTLLMMLARejectsDcpMultiTokenQuery(CustomTestCase):
         self.assertEqual(len(calls), 1)
 
 
+@unittest.skipUnless(
+    torch.cuda.is_available(), "the page-table build is a Triton kernel"
+)
+class TestDcpBlockTableIdSpace(CustomTestCase):
+    """`_fill_dcp_block_kv_indices` must emit entries in the pool's OWN id space.
+
+    On a static pool the DCP-collapsed page is already physical. Under the
+    unified memory pool it is still VIRTUAL, and an entry that skips the v2p
+    gather names whatever page happens to sit there now -- silent wrong-KV
+    reads rather than a crash, which is why the reference below is computed
+    from the tables rather than from the kernel.
+    """
+
+    PAGE_SIZE = 64
+    DCP_SIZE = 4
+    DCP_RANK = 2
+    MULTIPLIER = 3
+    # Virtual page per request, deliberately not the identity so a missing
+    # gather cannot coincide with the right answer.
+    VIRTUAL_PAGES = [[5, 2, 9], [7, 0, 4], [1, 8, 6]]
+    # Global KV lengths: one spanning 3 pages, one 2, one under a page.
+    SEQ_LENS = [3 * PAGE_SIZE * DCP_SIZE, 2 * PAGE_SIZE * DCP_SIZE - 1, 10]
+
+    def _make_backend(self, translator):
+        bs = len(self.VIRTUAL_PAGES)
+        max_pos = max(self.SEQ_LENS)
+        req_to_token = torch.full((bs, max_pos), -1, dtype=torch.int32, device="cuda")
+        span = self.PAGE_SIZE * self.DCP_SIZE
+        for req, pages in enumerate(self.VIRTUAL_PAGES):
+            for slot, virtual_page in enumerate(pages):
+                base = virtual_page * span
+                start = slot * span
+                end = min(start + span, self.SEQ_LENS[req])
+                if end <= start:
+                    break
+                req_to_token[req, start:end] = torch.arange(
+                    base, base + (end - start), dtype=torch.int32, device="cuda"
+                )
+        backend = object.__new__(TRTLLMMLABackend)
+        backend.page_size = self.PAGE_SIZE
+        backend.req_to_token = req_to_token
+        backend.kv_index_translator = translator
+        return backend, req_to_token
+
+    def _fill(self, translator):
+        backend, req_to_token = self._make_backend(translator)
+        bs = len(self.VIRTUAL_PAGES)
+        seq_lens = torch.tensor(self.SEQ_LENS, dtype=torch.int32, device="cuda")
+        local_seq_lens = get_dcp_lens(seq_lens, self.DCP_SIZE, self.DCP_RANK).to(
+            torch.int32
+        )
+        # One 128-page row: the padded width `_calc_padded_blocks` produces.
+        block_kv_indices = torch.full((bs, 128), -1, dtype=torch.int32, device="cuda")
+        parallel = SimpleNamespace(
+            dcp_enabled=True, dcp_size=self.DCP_SIZE, dcp_rank=self.DCP_RANK
+        )
+        with patch.object(backend_module, "get_parallel", return_value=parallel):
+            backend._fill_dcp_block_kv_indices(
+                block_kv_indices,
+                torch.arange(bs, dtype=torch.int64, device="cuda"),
+                local_seq_lens,
+            )
+        return block_kv_indices.cpu(), local_seq_lens.cpu(), req_to_token.cpu()
+
+    def _reference(self, req_to_token, local_seq_lens, v2p, multiplier):
+        """What each live entry must be, derived from the id-space definition."""
+        rows = []
+        for req in range(local_seq_lens.numel()):
+            local_pages = -(-int(local_seq_lens[req]) // self.PAGE_SIZE)
+            row = []
+            for page in range(local_pages):
+                pos = self.DCP_RANK + page * self.PAGE_SIZE * self.DCP_SIZE
+                widened = int(req_to_token[req, pos])
+                collapsed_page = widened // self.DCP_SIZE // self.PAGE_SIZE
+                if v2p is None:
+                    row.append(collapsed_page)
+                else:
+                    row.append(max(int(v2p[collapsed_page]) * multiplier, 0))
+            rows.append(row)
+        return rows
+
+    def _assert_matches(self, table, rows):
+        for req, row in enumerate(rows):
+            self.assertEqual(table[req, : len(row)].tolist(), row)
+            # Past the live prefix nothing is written: the caller-owned
+            # capture-stable buffer keeps what it had there.
+            self.assertTrue((table[req, len(row) :] == -1).all())
+
+    def test_static_pool_entries_are_the_collapsed_page(self):
+        translator = SimpleNamespace(full_v2p_table=None, full_page_multiplier=1)
+        table, local_lens, req_to_token = self._fill(translator)
+        self._assert_matches(table, self._reference(req_to_token, local_lens, None, 1))
+
+    def test_unified_pool_entries_go_through_the_page_table(self):
+        num_pages = 1 + max(max(p) for p in self.VIRTUAL_PAGES)
+        # A scrambled v2p, so an entry that skipped the gather would differ.
+        v2p = torch.tensor(
+            [(7 * i + 3) % num_pages for i in range(num_pages)],
+            dtype=torch.int64,
+            device="cuda",
+        )
+        translator = SimpleNamespace(
+            full_v2p_table=v2p, full_page_multiplier=self.MULTIPLIER
+        )
+        table, local_lens, req_to_token = self._fill(translator)
+        expected = self._reference(req_to_token, local_lens, v2p.cpu(), self.MULTIPLIER)
+        self._assert_matches(table, expected)
+        # And it is genuinely a translation, not an accident of the fixture.
+        static = self._reference(req_to_token, local_lens, None, 1)
+        self.assertNotEqual(expected, static)
+
+    def test_freed_page_lands_on_the_padding_sink(self):
+        # A tombstoned (-1) v2p row must clamp to entry 0, the reserved
+        # padding page, rather than scale -1 into a wild block id.
+        num_pages = 1 + max(max(p) for p in self.VIRTUAL_PAGES)
+        v2p = torch.arange(num_pages, dtype=torch.int64, device="cuda")
+        v2p[self.VIRTUAL_PAGES[0][0]] = -1
+        translator = SimpleNamespace(
+            full_v2p_table=v2p, full_page_multiplier=self.MULTIPLIER
+        )
+        table, _, _ = self._fill(translator)
+        self.assertEqual(int(table[0, 0]), 0)
+
+
+class TestFusedFp8WriteGate(CustomTestCase):
+    """The fused fp8 KV write must not serve a DCP-resolved loc.
+
+    Its only skip is the DCP owner rule (`vloc % world != rank`). Under the
+    unified pool that rule is already applied before the loc arrives, with
+    non-owned rows folded onto kernel id 0 -- so the kernel cannot skip the
+    padding sink the way `set_mla_kv_buffer`'s `reserved_skip_index` does, and
+    every non-owned token would store into slot 0.
+
+    A revert of the gate leaves accuracy tests green (slot 0 holds no real
+    data), which is why this is asserted directly.
+    """
+
+    def _gate(self, *, dcp_enabled: bool, is_translating: bool) -> bool:
+        backend = object.__new__(TRTLLMMLABackend)
+        backend.data_type = torch.float8_e4m3fn
+        backend.kv_lora_rank = 512
+        backend.qk_rope_head_dim = 64
+        backend.kv_index_translator = SimpleNamespace(is_translating=is_translating)
+        parallel = SimpleNamespace(
+            dcp_enabled=dcp_enabled,
+            attn_dcp_size=DCP_SIZE if dcp_enabled else 1,
+            attn_dcp_rank=DCP_RANK if dcp_enabled else 0,
+        )
+        with (
+            patch.object(backend_module, "get_parallel", return_value=parallel),
+            patch.object(
+                backend_module, "can_use_set_mla_kv_concat_q_fp8", return_value=True
+            ),
+            patch.object(
+                backend_module.envs.SGLANG_ENABLE_ASYNC_ASSERT, "get", lambda: False
+            ),
+        ):
+            return bool(
+                backend.data_type == torch.float8_e4m3fn
+                and not backend_module.envs.SGLANG_ENABLE_ASYNC_ASSERT.get()
+                and backend.kv_lora_rank == 512
+                and backend.qk_rope_head_dim == 64
+                and not (
+                    backend_module.get_parallel().dcp_enabled
+                    and backend.kv_index_translator.is_translating
+                )
+                and backend_module.can_use_set_mla_kv_concat_q_fp8()
+            )
+
+    def test_off_for_the_unified_pool_under_dcp(self):
+        self.assertFalse(self._gate(dcp_enabled=True, is_translating=True))
+
+    def test_on_for_every_other_combination(self):
+        # The gate must not cost the static pool or a non-DCP unified run
+        # their fused write.
+        self.assertTrue(self._gate(dcp_enabled=True, is_translating=False))
+        self.assertTrue(self._gate(dcp_enabled=False, is_translating=True))
+        self.assertTrue(self._gate(dcp_enabled=False, is_translating=False))
+
+    @unittest.skipUnless(torch.cuda.is_available(), "backend construction")
+    def test_helper_refuses_a_resolved_loc(self):
+        """Belt and braces: reaching the helper with a resolved loc asserts
+        rather than silently storing into the sink."""
+        backend = object.__new__(TRTLLMMLABackend)
+        backend.kv_index_translator = SimpleNamespace(is_translating=True)
+        parallel = SimpleNamespace(
+            dcp_enabled=True, attn_dcp_size=DCP_SIZE, attn_dcp_rank=DCP_RANK
+        )
+        layer = SimpleNamespace(
+            tp_q_head_num=1, v_head_dim=512, head_dim=576, layer_id=0
+        )
+        n = 4
+        with (
+            patch.object(backend_module, "get_parallel", return_value=parallel),
+            patch.object(
+                backend_module, "set_mla_kv_concat_q_fp8_covered", return_value=True
+            ),
+            patch.object(
+                backend,
+                "token_to_kv_pool",
+                SimpleNamespace(
+                    get_key_buffer=lambda _: torch.zeros(
+                        (8, 576), dtype=torch.uint8, device="cuda"
+                    )
+                ),
+                create=True,
+            ),
+        ):
+            with self.assertRaises(AssertionError):
+                backend._set_kv_and_concat_q_fp8_fused(
+                    layer=layer,
+                    loc=torch.zeros(n, dtype=torch.int64, device="cuda"),
+                    q=torch.zeros((n, 512), dtype=torch.bfloat16, device="cuda"),
+                    q_rope=torch.zeros((n, 64), dtype=torch.bfloat16, device="cuda"),
+                    k=torch.zeros((n, 512), dtype=torch.bfloat16, device="cuda"),
+                    k_rope=torch.zeros((n, 64), dtype=torch.bfloat16, device="cuda"),
+                )
+
+
 class TestDcpDecodeLayout(CustomTestCase):
     """Rank-local length math the decode page table above is built from."""
 

test/registered/models_e2e/test_kimi_linear_unified_memory_dcp_blackwell.py

@@ -0,0 +1,87 @@
+"""Kimi-Linear on the unified memory pool with decode context parallelism,
+served by the trtllm_mla family.
+
+Two read-index routes exist under `--enable-unified-memory --dcp-size > 1`.
+`test_kimi_linear_unified_memory.py` covers the flashinfer one, where
+`plan_dcp_decode_metadata` compacts this rank's ids and
+`translate_dcp_read_ids` converts them. This file covers the other: a block
+table built by `create_mla_kv_page_table_for_dcp`, which gathers a rank's
+cyclic slice and must take the pool's virtual->physical page table on the way.
+An entry that skipped that hop names whichever page sits there now, so the
+failure is stale KV rather than a crash -- which is what the GSM8K bar and the
+prefix-cache replay below are here to catch.
+
+`cutedsl_mla` stands in for the family: it is the DCP-native MLA decode kernel
+on Blackwell (the only backend flashinfer accepts `enable_dcp=True` for) and
+what the B300 hosts run, and prefill resolves to `trtllm_mla`, so one server
+covers both halves of the page-table contract. `trtllm_mla` and
+`tokenspeed_mla` share that builder and were checked by hand against the
+static pool (see below) rather than given a CI server each.
+
+Blackwell-only, hence its own module: the Hopper suite next door cannot host
+it. In extra-b rather than nightly because what it guards is a per-PR argument
+gate, which a regression closes at boot rather than degrading overnight.
+
+Measured on B300 (1 sigma ~= 0.02). This file scores 0.920 and takes 2.5 min.
+Matched unified-vs-static pairs, run by hand at 200 questions, to show the two
+pools agree:
+
+                                    unified / static
+    TP2 DCP2  cutedsl_mla           0.905 / 0.905
+    TP2 DCP2  trtllm_mla            0.895 / 0.895
+    TP2 DCP2  tokenspeed_mla        0.890 / 0.895, repeat 0.890 / 0.885
+    TP4 DCP4  cutedsl_mla           0.910 / 0.900
+
+The cutedsl_mla and trtllm_mla pairs reproduced exactly. tokenspeed_mla moved
+on a repeat of the same questions on BOTH pools, so its spread is the backend,
+not the pool. The TP4 DCP4 row is there because it widens the virtual page to
+4 x 64 = 256, past what this file exercises.
+
+    python -m pytest test/registered/models_e2e/test_kimi_linear_unified_memory_dcp_blackwell.py -v
+"""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
+from sglang.test.server_fixtures.default_fixture import DefaultServerBase
+
+register_cuda_ci(est_time=250, stage="extra-b", runner_config="4-gpu-b200")
+
+KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
+
+
+class TestKimiLinearUnifiedMemoryDCPCuteDsl(
+    GSM8KMixin, PrefixCacheBranchingMixin, DefaultServerBase
+):
+    """cutedsl_mla decode, trtllm_mla prefill (what cutedsl_mla resolves to)."""
+
+    model = KIMI_LINEAR_MODEL
+    cache_chunk_size = 64
+    gsm8k_score_threshold = 0.88
+    other_args = [
+        "--trust-remote-code",
+        "--tp-size",
+        "2",
+        "--dcp-size",
+        "2",
+        "--attention-backend",
+        "cutedsl_mla",
+        "--chunked-prefill-size",
+        "2048",
+        # The per-batch DCP gather buffer is sized by the batch's total KV, so
+        # the default static fraction leaves a 200-question burst ~1 GB and it
+        # OOMs inside an NCCL collective.
+        "--mem-fraction-static",
+        "0.80",
+        "--max-running-requests",
+        "128",
+        "--cuda-graph-max-bs-decode",
+        "128",
+        "--enable-unified-memory",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()