[Feature] Unified memory: support decode context parallelism for Kimi-Linear (#36890)

2026-09-01 12:44:26 -07:00
parent 3315356cc0
commit 0b1ce3d140
20 changed files with 671 additions and 105 deletions
@@ -70,7 +70,9 @@ class TestFilterDcpLocalChunkKvIndices(CustomTestCase):
        return torch.cat(runs) if runs else torch.empty(0, dtype=torch.int64)

    def _owner_rule(self, kv, dcp_size, dcp_rank):
-        return kv[kv % dcp_size == dcp_rank] // dcp_size
+        # Selection only: the filters leave ids WIDENED and the collapse now
+        # happens once, in KVIndexTranslator.translate_dcp_read_ids.
+        return kv[kv % dcp_size == dcp_rank]

    def _run(self, starts, lens, dcp_size, dcp_rank, seed=0):
        kv = self._build_chunk(starts, lens, dcp_size, seed)