chore: import upstream snapshot with attribution

python/tokenspeed/runtime/models/base/decoder_layer.py

@@ -0,0 +1,369 @@
+# 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.
+
+"""Base decoder layer classes.
+
+``BaseDecoderLayer`` uses CommManager for communication (the default path).
+``CompiledDecoderLayer`` uses the compiler-driven path.
+"""
+
+from __future__ import annotations
+
+from typing import Generic, TypeVar
+
+import torch
+from torch import nn
+from transformers import PretrainedConfig
+
+from tokenspeed.runtime.distributed.comm_manager import CommManager
+from tokenspeed.runtime.distributed.mapping import Mapping
+from tokenspeed.runtime.execution.context import ForwardContext
+from tokenspeed.runtime.layers.layernorm import RMSNorm
+from tokenspeed.runtime.layers.quantization import QuantizationConfig as Q
+from tokenspeed.runtime.models.base.execution import (
+    CompiledDecoderLayer as _CompiledRuntime,
+)
+from tokenspeed.runtime.models.base.execution import (
+    ExecutionNode,
+)
+from tokenspeed.runtime.models.base.module_spec import ModuleKind, ModuleSpec
+from tokenspeed.runtime.models.base.placement import ParallelGroup, Partial, Replicate
+
+
+def _default_compute_output_placement(
+    mapping: Mapping,
+    group: ParallelGroup,
+) -> Partial | None:
+    if group == ParallelGroup.ATTN_TP:
+        has_parallel = mapping.has_attn_tp
+    elif group == ParallelGroup.DENSE_TP:
+        has_parallel = mapping.dense.has_tp
+    elif group == ParallelGroup.MOE_TP_EP:
+        has_parallel = mapping.moe.has_tp_ep
+    else:
+        raise ValueError(f"Unknown group: {group}")
+    return Partial(group) if has_parallel else None
+
+
+_C = TypeVar("_C", bound=PretrainedConfig)
+
+
+class BaseDecoderLayer(nn.Module, Generic[_C]):
+    """Default decoder layer using CommManager for communication.
+
+    Subclasses override ``resolve_attn()`` and ``resolve_mlp()``.
+    """
+
+    def __init__(
+        self,
+        config: _C,
+        layer_id: int,
+        mapping: Mapping,
+        quant_config: Q | None = None,
+        prefix: str = "",
+    ) -> None:
+
+        super().__init__()
+
+        self.config = config
+        self.quant_config = quant_config
+        self.layer_id = layer_id
+        self.total_layers = config.num_hidden_layers
+        self.mapping = mapping
+
+        self.input_layernorm = self.resolve_norm()
+        self.post_attention_layernorm = self.resolve_norm()
+
+        self.self_attn = self.resolve_attn(prefix)
+        self.mlp = self.resolve_mlp(prefix)
+
+        self.comm_manager = CommManager(
+            mapping=self.mapping,
+            layer_id=layer_id,
+            is_moe=self.is_moe_layer,
+            prev_is_moe=self.is_moe_layer,
+            input_layernorm=self.input_layernorm,
+            post_attn_layernorm=self.post_attention_layernorm,
+        )
+
+    @property
+    def is_moe_layer(self) -> bool:
+
+        return False
+
+    def resolve_norm(self) -> nn.Module:
+
+        return RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
+
+    def resolve_attn(self, prefix: str) -> nn.Module:
+
+        raise NotImplementedError
+
+    def resolve_mlp(self, prefix: str) -> nn.Module:
+
+        raise NotImplementedError
+
+    def forward_attn(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        ctx: ForwardContext,
+        out_cache_loc: torch.Tensor,
+        residual: torch.Tensor | None,
+        aux_hidden_states: list | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        hidden_states, residual = self.comm_manager.input_reduce_norm(
+            hidden_states, residual
+        )
+
+        if aux_hidden_states is not None:
+            # Under RSAG the residual entering this layer is reduce-scattered
+            # across the attn TP group; aux consumers (e.g. the EAGLE3
+            # drafter) expect full rows, so gather before capturing.
+            aux_hidden_states.append(
+                self.comm_manager.gather_residual(residual, ctx).clone()
+            )
+
+        hidden_states = self.comm_manager.pre_attn_comm(hidden_states, ctx)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            ctx=ctx,
+            out_cache_loc=out_cache_loc,
+        )
+
+        hidden_states, residual = self.comm_manager.post_attn_reduce_norm(
+            hidden_states, residual, ctx
+        )
+
+        return hidden_states, residual
+
+    def forward_mlp(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        ctx: ForwardContext,
+        num_global_tokens: int,
+        max_num_tokens_per_gpu: int,
+    ) -> torch.Tensor:
+
+        hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx)
+
+        if self.is_moe_layer:
+
+            hidden_states = self.mlp(
+                hidden_states, num_global_tokens, max_num_tokens_per_gpu
+            )
+
+        else:
+
+            hidden_states = self.mlp(hidden_states)
+
+        hidden_states, residual = self.comm_manager.post_mlp_fused(
+            hidden_states, residual, ctx
+        )
+
+        return hidden_states
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        ctx: ForwardContext,
+        out_cache_loc: torch.Tensor,
+        residual: torch.Tensor | None,
+        aux_hidden_states: list | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        num_global_tokens, max_num_tokens_per_gpu = self.comm_manager.get_num_tokens(
+            ctx
+        )
+
+        if not ctx.forward_mode.is_idle():
+
+            hidden_states, residual = self.forward_attn(
+                positions,
+                hidden_states,
+                ctx,
+                out_cache_loc,
+                residual,
+                aux_hidden_states,
+            )
+
+            hidden_states = self.forward_mlp(
+                hidden_states,
+                residual,
+                ctx,
+                num_global_tokens,
+                max_num_tokens_per_gpu,
+            )
+
+        else:
+
+            hidden_states = self.forward_mlp(
+                hidden_states,
+                residual,
+                ctx,
+                num_global_tokens,
+                max_num_tokens_per_gpu,
+            )
+
+        return hidden_states, residual
+
+
+class BaseMoEDecoderLayer(BaseDecoderLayer):
+
+    @property
+    def is_moe_layer(self) -> bool:
+
+        return True
+
+
+class CompiledDecoderLayer(nn.Module, Generic[_C]):
+    """Compiler-driven decoder layer (opt-in).
+
+    Instead of CommManager, the forward delegates to a
+    ``_CompiledRuntime`` produced by the layer compiler.
+    """
+
+    def __init__(
+        self,
+        config: _C,
+        layer_id: int,
+        mapping: Mapping,
+        quant_config: Q | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        self.quant_config = quant_config
+        self.layer_id = layer_id
+        self.total_layers = config.num_hidden_layers
+        self.mapping = mapping
+        self.prefix = prefix
+
+        self._compiled: _CompiledRuntime | None = None
+        self._exec_plan = self.build_execution_plan(prefix)
+
+    @property
+    def is_moe_layer(self) -> bool:
+        return False
+
+    def resolve_norm(self) -> nn.Module:
+        return RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
+
+    def build_execution_plan(self, prefix: str) -> list[ExecutionNode]:
+        self.input_layernorm = self.resolve_norm()
+        self.self_attn = self.resolve_attn(prefix)
+        self.post_attention_layernorm = self.resolve_norm()
+        self.mlp = self.resolve_mlp(prefix)
+
+        return [
+            ExecutionNode(
+                module=self.input_layernorm,
+                spec=self.norm_spec(captures_aux=True, skip_on_idle=True),
+                name="input_layernorm",
+            ),
+            ExecutionNode(
+                module=self.self_attn,
+                spec=self.attn_spec(),
+                name="self_attn",
+            ),
+            ExecutionNode(
+                module=self.post_attention_layernorm,
+                spec=self.norm_spec(),
+                name="post_attention_layernorm",
+            ),
+            ExecutionNode(
+                module=self.mlp,
+                spec=self.mlp_spec(),
+                name="mlp",
+            ),
+        ]
+
+    def norm_spec(
+        self,
+        *,
+        captures_aux: bool = False,
+        skip_on_idle: bool = False,
+    ) -> ModuleSpec:
+        return ModuleSpec.from_kind(
+            kind=ModuleKind.NORM,
+            supports_fused_reduce_norm=True,
+            captures_aux=captures_aux,
+            skip_on_idle=skip_on_idle,
+        )
+
+    def attn_spec(self) -> ModuleSpec:
+        input_placement = Replicate(ParallelGroup.ATTN_TP)
+        return ModuleSpec.from_kind(
+            input_placement=input_placement,
+            output_placement=_default_compute_output_placement(
+                self.mapping, ParallelGroup.ATTN_TP
+            ),
+            kind=ModuleKind.ATTENTION,
+            skip_on_idle=True,
+        )
+
+    def mlp_spec(self) -> ModuleSpec:
+        mlp_group = (
+            ParallelGroup.MOE_TP_EP if self.is_moe_layer else ParallelGroup.DENSE_TP
+        )
+        kind = ModuleKind.MOE if self.is_moe_layer else ModuleKind.DENSE_MLP
+        return ModuleSpec.from_kind(
+            input_placement=Replicate(mlp_group),
+            output_placement=_default_compute_output_placement(self.mapping, mlp_group),
+            kind=kind,
+        )
+
+    def resolve_attn(self, prefix: str) -> nn.Module:
+        raise NotImplementedError
+
+    def resolve_mlp(self, prefix: str) -> nn.Module:
+        raise NotImplementedError
+
+    def resolve_exec_plan(self) -> list[ExecutionNode]:
+        return self._exec_plan
+
+    def set_compiled(self, compiled: _CompiledRuntime) -> None:
+        self._compiled = compiled
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        ctx: ForwardContext,
+        out_cache_loc: torch.Tensor,
+        residual: torch.Tensor | None,
+        aux_hidden_states: list | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self._compiled.forward(
+            positions, hidden_states, ctx, out_cache_loc, residual, aux_hidden_states
+        )
+
+
+class CompiledMoEDecoderLayer(CompiledDecoderLayer):
+
+    @property
+    def is_moe_layer(self) -> bool:
+        return True