[CP] Consolidate decode-context-parallel (DCP) helpers into layers/dcp/ (#29365)

Co-authored-by: Hao Phan <htphan@nvidia.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

test/registered/dcp/test_dcp_layout_unit.py

@@ -0,0 +1,91 @@
+"""CPU unit test for the decode-context-parallel (DCP) per-rank KV-length math.
+
+Pins ``get_dcp_lens`` (the single, superset implementation in
+``layers/dcp/layout.py``) to a brute-force owner-count reference, and proves
+it is bit-identical to the legacy in-place formula that
+``update_local_kv_lens_for_dcp`` used before it was collapsed into a wrapper:
+
+    floor((len - rank - 1) / N) + 1   ==   len // N + (rank < len % N)   (len >= 0)
+
+Usage:
+    python -m pytest test_dcp_layout_unit.py -v
+    python test_dcp_layout_unit.py
+"""
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.dcp.layout import get_dcp_lens
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+DCP_SIZES = [1, 2, 3, 4, 8]
+LENS = list(range(0, 41))
+STARTS = [0, 1, 2, 5, 7, 13, 31]
+
+
+def _owner_count(length: int, n: int, rank: int, start: int) -> int:
+    """Ground truth: # of absolute positions p in [start, start+length) with p % n == rank."""
+    return sum(1 for p in range(start, start + length) if p % n == rank)
+
+
+def _legacy_inplace_formula(length: int, n: int, rank: int) -> int:
+    """The pre-refactor update_local_kv_lens_for_dcp body (start == 0 case)."""
+    return (length - rank - 1) // n + 1
+
+
+class TestGetDcpLens(unittest.TestCase):
+    def test_start_none_matches_owner_count(self):
+        for n in DCP_SIZES:
+            for rank in range(n):
+                lens = torch.tensor(LENS, dtype=torch.int32)
+                got = get_dcp_lens(lens, n, rank)
+                expected = torch.tensor(
+                    [_owner_count(L, n, rank, 0) for L in LENS], dtype=torch.int32
+                )
+                self.assertTrue(
+                    torch.equal(got.to(torch.int32), expected),
+                    f"start=None mismatch at n={n}, rank={rank}: {got.tolist()} != {expected.tolist()}",
+                )
+
+    def test_start_none_matches_legacy_inplace_formula(self):
+        # The collapse claim: get_dcp_lens (start=None) == legacy floor((L-rank-1)/N)+1.
+        for n in DCP_SIZES:
+            for rank in range(n):
+                lens = torch.tensor(LENS, dtype=torch.int64)
+                got = get_dcp_lens(lens, n, rank)
+                legacy = torch.tensor(
+                    [_legacy_inplace_formula(L, n, rank) for L in LENS],
+                    dtype=torch.int64,
+                )
+                self.assertTrue(
+                    torch.equal(got.to(torch.int64), legacy),
+                    f"legacy-formula mismatch at n={n}, rank={rank}",
+                )
+
+    def test_start_tensor_matches_owner_count(self):
+        for n in DCP_SIZES:
+            for rank in range(n):
+                for start in STARTS:
+                    lens = torch.tensor(LENS, dtype=torch.int64)
+                    start_t = torch.full_like(lens, start)
+                    got = get_dcp_lens(lens, n, rank, start=start_t)
+                    expected = torch.tensor(
+                        [_owner_count(L, n, rank, start) for L in LENS],
+                        dtype=torch.int64,
+                    )
+                    self.assertTrue(
+                        torch.equal(got.to(torch.int64), expected),
+                        f"start={start} mismatch at n={n}, rank={rank}: "
+                        f"{got.tolist()} != {expected.tolist()}",
+                    )
+
+    def test_dcp_size_one_is_identity(self):
+        lens = torch.tensor(LENS, dtype=torch.int32)
+        self.assertTrue(torch.equal(get_dcp_lens(lens, 1, 0), lens))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/dcp/test_dsv31_dcp8_gsm8k.py

@@ -126,7 +126,7 @@ class TestDSV31DCP8TP8GSM8K(GSM8KMixin, BasicDecodeCorrectnessMixin, CustomTestC
 
     This test exercises the full DCP decode and extend paths:
       - Decode: query all-gather → attention on local KV shard → LSE
-        correction via cp_lse_ag_out_rs → reduce-scatter
+        correction via cp_lse_ag_out_rs_mla → reduce-scatter
       - Extend (prefill): all-gather prefix KV cache across DCP ranks,
         attend with full context
 
@@ -205,7 +205,7 @@ class TestDSV31DCP8LogprobParity(BasicDecodeCorrectnessMixin, CustomTestCase):
            introduces small numerical differences)
 
     This catches subtle correctness bugs in the DCP LSE correction path
-    (cp_lse_ag_out_rs) that a coarse GSM8K accuracy gate cannot detect.
+    (cp_lse_ag_out_rs_mla) that a coarse GSM8K accuracy gate cannot detect.
     For example, if exp2/exp mismatch causes a systematic bias in the
     attention output, logprobs will diverge by more than the tolerance.
     """