[Perf] Unified memory: close the DCP decode gap on Blackwell (#37926)

test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py

@@ -432,6 +432,89 @@ class TestFusedFp8WriteGate(CustomTestCase):
                 )
 
 
+@unittest.skipUnless(torch.cuda.is_available(), "the fused translate is Triton")
+class TestFusedWriteLocTranslateCuda(CustomTestCase):
+    """The fused write-loc translate must agree with its own CPU branch.
+
+    The two implementations must not drift: Triton truncates division toward
+    zero where torch floors it, so a negative loc and a tombstoned v2p row are
+    where a divergence would appear -- and the CPU branch is all the CPU suites
+    ever exercise.
+    """
+
+    def test_cuda_matches_the_cpu_branch(self):
+        from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids
+
+        for page_size in (1, 64):
+            span = page_size * 4
+            loc = torch.tensor(
+                [-1, 0, 1, page_size, span, span + 1, 2 * span + 3, 5 * page_size],
+                dtype=torch.int64,
+            )
+            # Size the table past the highest page any loc can name, then
+            # scramble it and tombstone one row (-1), so neither a dropped
+            # clamp nor a skipped gather can coincide with the right answer.
+            num_pages = int(loc.max()) // page_size + 2
+            v2p = torch.tensor(
+                [(5 * i + 2) % num_pages for i in range(num_pages)] + [-1],
+                dtype=torch.int64,
+            )
+            v2p[1] = -1
+            for dcp_size, dcp_rank in ((1, 0), (2, 1), (4, 2)):
+                kw = dict(
+                    page_size=page_size,
+                    stride=page_size * 3,
+                    dcp_size=dcp_size,
+                    dcp_rank=dcp_rank,
+                )
+                cpu = write_loc_to_kernel_ids(loc=loc, v2p=v2p, **kw)
+                gpu = write_loc_to_kernel_ids(loc=loc.cuda(), v2p=v2p.cuda(), **kw)
+                self.assertEqual(
+                    cpu.tolist(),
+                    gpu.cpu().tolist(),
+                    f"ps={page_size} dcp_size={dcp_size} rank={dcp_rank}",
+                )
+
+    def test_wide_out_clears_the_stale_tail(self):
+        """`out_width` past the batch must zero the tail in the same launch.
+
+        This is what lets a backend hand in its whole capture-stable buffer:
+        a shorter replay leaves stale kernel-facing ids past the batch, and the
+        captured write kernel consumes the full buffer, so an uncleared tail
+        scatters pad rows into live KV pages.
+        """
+        from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids
+
+        v2p = torch.tensor([2, 5, 1, 3, 4], dtype=torch.int64, device="cuda")
+        loc = torch.tensor([0, 64, 128], dtype=torch.int64, device="cuda")
+        width = 8
+        # Poison the whole buffer so an unwritten or uncleared cell is visible.
+        buf = torch.full((width,), -999, dtype=torch.int64, device="cuda")
+        write_loc_to_kernel_ids(
+            loc=loc, v2p=v2p, page_size=64, stride=64 * 2, out=buf, out_width=width
+        )
+        self.assertEqual(buf[:3].tolist(), [2 * 128, 5 * 128, 1 * 128])
+        self.assertEqual(buf[3:].tolist(), [0] * (width - 3))
+
+        # And it must agree with the narrow call on the live prefix.
+        narrow = write_loc_to_kernel_ids(loc=loc, v2p=v2p, page_size=64, stride=64 * 2)
+        self.assertEqual(narrow.tolist(), buf[:3].tolist())
+
+    def test_out_is_written_in_place(self):
+        # The captured decode path hands in a capture-stable buffer; rebinding
+        # instead of filling it would leave the graph on a stale pointer.
+        from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids
+
+        v2p = torch.tensor([2, 5, -1, 3], dtype=torch.int64, device="cuda")
+        loc = torch.tensor([0, 64, 128, 192], dtype=torch.int64, device="cuda")
+        dst = torch.full_like(loc, -7)
+        ret = write_loc_to_kernel_ids(
+            loc=loc, v2p=v2p, page_size=64, stride=64 * 2, out=dst
+        )
+        self.assertIs(ret, dst)
+        self.assertEqual(dst.tolist(), [2 * 128, 5 * 128, 0, 3 * 128])
+
+
 class TestDcpDecodeLayout(CustomTestCase):
     """Rank-local length math the decode page table above is built from."""