chore: import upstream snapshot with attribution

python/tokenspeed/runtime/distributed/mapping.py

@@ -0,0 +1,392 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import math
+from functools import cached_property
+
+Group = tuple[int, ...]
+
+
+def _resolve_parallelism_sizes(world_size: int, *sizes: int | None) -> tuple[int, ...]:
+    """Resolve the parallelism sizes given world_size.
+
+    `sizes` is ordered innermost (fastest-varying) to outermost.
+    """
+    assert all(x is None or x > 0 for x in sizes)
+
+    resolved = [x for x in sizes]
+    num_to_resolve = sum(x is None for x in sizes)
+    if num_to_resolve > 0:
+        provided_size = math.prod(x for x in sizes if x is not None)
+        assert provided_size <= world_size
+        assert world_size % provided_size == 0
+        resolved_size = world_size // provided_size
+
+        for index, size in enumerate(resolved):
+            if size is None:
+                resolved[index] = resolved_size
+                resolved_size = 1
+
+    assert math.prod(resolved) == world_size
+    return tuple(resolved)
+
+
+def _make_parallelism_rank(rank: int, size: int, stride: int = 1) -> int:
+    """Return the rank of given size and stride."""
+    return (rank // stride) % size
+
+
+def _make_parallelism_group(rank: int, size: int, stride: int = 1) -> Group:
+    """Return the group of ranks of given size and stride."""
+    base = rank - (rank // stride % size) * stride
+    return tuple(base + j * stride for j in range(size))
+
+
+class MappingBase:
+
+    def __init__(self, rank: int | None = None, world_size: int = 1):
+        assert rank is None or rank >= 0
+        self._rank = rank
+        assert world_size > 0
+        self._world_size = world_size
+
+    @property
+    def rank(self) -> int:
+        assert self._rank is not None, "rank is not initialized"
+        return self._rank
+
+    @rank.setter
+    def rank(self, rank: int):
+        assert self._rank is None, "rank is already initialized"
+        assert rank >= 0
+        self._rank = rank
+        self._on_rank_initialized(rank)
+
+    def _on_rank_initialized(self, rank: int):
+        return None
+
+    @property
+    def world_size(self) -> int:
+        return self._world_size
+
+    @cached_property
+    def world_group(self) -> Group:
+        return _make_parallelism_group(self.rank, self.world_size, stride=1)
+
+
+class DenseLayerMapping(MappingBase):
+
+    def __init__(
+        self,
+        rank: int | None = None,
+        world_size: int = 1,
+        tp_size: int | None = None,
+        dp_size: int | None = None,
+    ):
+        super().__init__(rank, world_size)
+        self.tp_size, self.dp_size = _resolve_parallelism_sizes(
+            self.world_size, tp_size, dp_size
+        )
+
+    @cached_property
+    def has_tp(self) -> bool:
+        return self.tp_size > 1
+
+    @cached_property
+    def tp_rank(self) -> int:
+        return _make_parallelism_rank(self.rank, self.tp_size, stride=1)
+
+    @cached_property
+    def tp_group(self) -> Group:
+        return _make_parallelism_group(self.rank, self.tp_size, stride=1)
+
+    @cached_property
+    def has_dp(self) -> bool:
+        return self.dp_size > 1
+
+    @cached_property
+    def dp_rank(self) -> int:
+        return _make_parallelism_rank(self.rank, self.dp_size, stride=self.tp_size)
+
+    @cached_property
+    def dp_group(self) -> Group:
+        return _make_parallelism_group(self.rank, self.dp_size, stride=self.tp_size)
+
+
+class AttentionLayerMapping(MappingBase):
+
+    def __init__(
+        self,
+        rank: int | None = None,
+        world_size: int = 1,
+        tp_size: int | None = None,
+        cp_size: int | None = None,
+        dp_size: int | None = None,
+    ):
+        super().__init__(rank, world_size)
+        self.tp_size, self.cp_size, self.dp_size = _resolve_parallelism_sizes(
+            self.world_size, tp_size, cp_size, dp_size
+        )
+
+    @cached_property
+    def has_tp(self) -> bool:
+        return self.tp_size > 1
+
+    @cached_property
+    def tp_rank(self) -> int:
+        return _make_parallelism_rank(self.rank, self.tp_size, stride=1)
+
+    @cached_property
+    def tp_group(self) -> Group:
+        return _make_parallelism_group(self.rank, self.tp_size, stride=1)
+
+    @cached_property
+    def has_cp(self) -> bool:
+        return self.cp_size > 1
+
+    @cached_property
+    def cp_rank(self) -> int:
+        return _make_parallelism_rank(self.rank, self.cp_size, stride=self.tp_size)
+
+    @cached_property
+    def cp_group(self) -> Group:
+        return _make_parallelism_group(self.rank, self.cp_size, stride=self.tp_size)
+
+    @cached_property
+    def has_dp(self) -> bool:
+        return self.dp_size > 1
+
+    @cached_property
+    def dp_rank(self) -> int:
+        return _make_parallelism_rank(
+            self.rank, self.dp_size, stride=self.tp_size * self.cp_size
+        )
+
+    @cached_property
+    def dp_group(self) -> Group:
+        return _make_parallelism_group(
+            self.rank, self.dp_size, stride=self.tp_size * self.cp_size
+        )
+
+    def scatter_index(self, rank: int) -> int:
+        """Index of ``rank`` in a dp-major/tp-minor scattered token count
+        table; cp peers share their dp group's tp split."""
+        tp_rank = _make_parallelism_rank(rank, self.tp_size, stride=1)
+        dp_rank = _make_parallelism_rank(
+            rank, self.dp_size, stride=self.tp_size * self.cp_size
+        )
+        return dp_rank * self.tp_size + tp_rank
+
+
+class MoeLayerMapping(MappingBase):
+    def __init__(
+        self,
+        rank: int | None = None,
+        world_size: int = 1,
+        tp_size: int | None = None,
+        ep_size: int | None = None,
+        dp_size: int | None = None,
+    ):
+        super().__init__(rank, world_size)
+        self.tp_size, self.ep_size, self.dp_size = _resolve_parallelism_sizes(
+            self.world_size, tp_size, ep_size, dp_size
+        )
+
+    @cached_property
+    def has_tp(self) -> bool:
+        return self.tp_size > 1
+
+    @cached_property
+    def tp_rank(self) -> int:
+        return _make_parallelism_rank(self.rank, self.tp_size, stride=1)
+
+    @cached_property
+    def tp_group(self) -> Group:
+        return _make_parallelism_group(self.rank, self.tp_size, stride=1)
+
+    @cached_property
+    def has_ep(self) -> bool:
+        return self.ep_size > 1
+
+    @cached_property
+    def ep_rank(self) -> int:
+        return _make_parallelism_rank(self.rank, self.ep_size, stride=self.tp_size)
+
+    @cached_property
+    def ep_group(self) -> Group:
+        return _make_parallelism_group(self.rank, self.ep_size, stride=self.tp_size)
+
+    @cached_property
+    def has_tp_ep(self) -> bool:
+        return self.tp_ep_size > 1
+
+    @cached_property
+    def tp_ep_size(self) -> int:
+        return self.tp_size * self.ep_size
+
+    @cached_property
+    def tp_ep_rank(self) -> int:
+        return _make_parallelism_rank(self.rank, self.tp_ep_size, stride=1)
+
+    @cached_property
+    def tp_ep_group(self) -> Group:
+        return _make_parallelism_group(self.rank, self.tp_ep_size, stride=1)
+
+    @cached_property
+    def has_dp(self) -> bool:
+        return self.dp_size > 1
+
+    @cached_property
+    def dp_rank(self) -> int:
+        return _make_parallelism_rank(
+            self.rank, self.dp_size, stride=self.tp_size * self.ep_size
+        )
+
+    @cached_property
+    def dp_group(self) -> Group:
+        return _make_parallelism_group(
+            self.rank, self.dp_size, stride=self.tp_size * self.ep_size
+        )
+
+
+class VisionTowerMapping(MappingBase):
+    """Parallel mapping for vision encoders. Vision layers run colocated and
+    share the attention TP group; non-colocated deployments should run the
+    encoder out-of-engine (EPD-style workers + gateway dispatch).
+    """
+
+    def __init__(
+        self,
+        rank: int | None = None,
+        world_size: int = 1,
+        tp_size: int | None = None,
+    ):
+        super().__init__(rank, world_size)
+        (self.tp_size,) = _resolve_parallelism_sizes(self.world_size, tp_size)
+
+    @cached_property
+    def has_tp(self) -> bool:
+        return self.tp_size > 1
+
+    @cached_property
+    def tp_rank(self) -> int:
+        return _make_parallelism_rank(self.rank, self.tp_size, stride=1)
+
+    @cached_property
+    def tp_group(self) -> Group:
+        return _make_parallelism_group(self.rank, self.tp_size, stride=1)
+
+
+class Mapping(MappingBase):
+
+    def __init__(
+        self,
+        rank: int | None = None,
+        world_size: int = 1,
+        *,
+        attn_tp_size: int | None = None,
+        attn_cp_size: int | None = None,
+        attn_dp_size: int | None = None,
+        dense_tp_size: int | None = None,
+        dense_dp_size: int | None = None,
+        moe_tp_size: int | None = None,
+        moe_ep_size: int | None = None,
+        moe_dp_size: int | None = None,
+        nprocs_per_node: int | None = None,
+        nnodes: int | None = None,
+        base_gpu_id: int = 0,
+        gpu_id_step: int = 1,
+    ):
+        super().__init__(rank, world_size)
+        self.attn = AttentionLayerMapping(
+            rank=rank,
+            world_size=world_size,
+            tp_size=attn_tp_size,
+            cp_size=attn_cp_size,
+            dp_size=attn_dp_size,
+        )
+        self.dense = DenseLayerMapping(
+            rank=rank,
+            world_size=world_size,
+            tp_size=dense_tp_size,
+            dp_size=dense_dp_size,
+        )
+        self.moe = MoeLayerMapping(
+            rank=rank,
+            world_size=world_size,
+            tp_size=moe_tp_size,
+            ep_size=moe_ep_size,
+            dp_size=moe_dp_size,
+        )
+        # Vision tower runs colocated on the attention TP group.
+        self.vision = VisionTowerMapping(
+            rank=rank,
+            world_size=self.attn.tp_size,
+            tp_size=self.attn.tp_size,
+        )
+        self.nprocs_per_node, self.nnodes = _resolve_parallelism_sizes(
+            self.world_size, nprocs_per_node, nnodes
+        )
+        assert base_gpu_id >= 0
+        assert gpu_id_step > 0
+        self.base_gpu_id = base_gpu_id
+        self.gpu_id_step = gpu_id_step
+
+    def _on_rank_initialized(self, rank: int):
+        self.attn.rank = rank
+        self.dense.rank = rank
+        self.moe.rank = rank
+        self.vision.rank = rank
+
+    @cached_property
+    def has_attn_tp(self) -> bool:
+        return self.attn.has_tp
+
+    @cached_property
+    def has_attn_cp(self) -> bool:
+        return self.attn.has_cp
+
+    @cached_property
+    def has_attn_dp(self) -> bool:
+        return self.attn.has_dp
+
+    @cached_property
+    def node_rank(self) -> int:
+        return self.rank // self.nprocs_per_node
+
+    @cached_property
+    def local_rank(self) -> int:
+        return self.rank % self.nprocs_per_node
+
+    @cached_property
+    def gpu_id(self) -> int:
+        return self.base_gpu_id + self.local_rank * self.gpu_id_step
+
+    def __repr__(self) -> str:
+        rank_str = str(self._rank) if self._rank is not None else "?"
+        lines = [
+            f"Mapping(rank={rank_str}, world_size={self.world_size})",
+            f"  Cluster : {self.nnodes} node(s) x {self.nprocs_per_node} proc(s)",
+            f"  Attention: tp={self.attn.tp_size}  cp={self.attn.cp_size}  dp={self.attn.dp_size}",
+            f"    Vision: tp={self.vision.tp_size}",
+            f"  Dense   : tp={self.dense.tp_size}  dp={self.dense.dp_size}",
+            f"  MoE     : tp={self.moe.tp_size}  ep={self.moe.ep_size}  dp={self.moe.dp_size}",
+        ]
+        return "\n".join(lines)