RankMixer：工业级推荐排序模型的规模化之路

import torch
import torch.nn as nn
import torch.nn.functional as F


class FeatureTokenizer(nn.Module):
    """
    将拼接后的异构特征 embedding 切分为 T 个固定维度的 token，
    并投影到统一的模型宽度 D。
    对应论文公式 (2): x_i = Proj(e_input[d*(i-1):d*i])
    """

    def __init__(self, input_dim: int, num_tokens: int, model_dim: int):
        super().__init__()
        assert input_dim % num_tokens == 0, "特征分组需能被 T 整除，否则请先padding"
        self.num_tokens = num_tokens
        self.chunk_dim = input_dim // num_tokens  # 对应论文中的 d
        self.model_dim = model_dim

        # 每个 token 有独立的投影矩阵（Proj 函数），
        # 也可以选择所有 token 共享一个 Proj，这里采用逐 token 独立投影
        self.proj = nn.ModuleList([
            nn.Linear(self.chunk_dim, model_dim) for _ in range(num_tokens)
        ])

    def forward(self, e_input: torch.Tensor) -> torch.Tensor:
        """
        e_input: [B, input_dim]  已经按语义分组拼接好的特征向量
        return:  [B, T, D]
        """
        chunks = e_input.split(self.chunk_dim, dim=-1)  # T 个 [B, chunk_dim]
        tokens = [proj(chunk) for proj, chunk in zip(self.proj, chunks)]
        return torch.stack(tokens, dim=1)  # [B, T, D]

class MultiHeadTokenMixing(nn.Module):
    """
    对应论文 3.3.1 节 & 公式 (3)(4)(5)。
    输入 X: [B, T, D] -> 输出 S: [B, T, D] （因为设定 H = T，token数不变）

    实现思路：
      1. 每个 token 切分为 H 个 head，每个 head 维度 D // H
      2. 按 head 维度重组：把所有 token 的第 h 个 head 拼在一起，形成新 token s^h
      3. 因为设定 H = T，输出 token 数与输入一致，天然可以做残差连接
    """

    def __init__(self, num_tokens: int, model_dim: int):
        super().__init__()
        self.T = num_tokens
        self.D = model_dim
        self.H = num_tokens  # 论文设定 H = T
        assert model_dim % self.H == 0, "D 必须能被 H 整除"
        self.head_dim = model_dim // self.H  # D // H

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: [B, T, D]
        return: [B, T, D]  (H个mixed token，每个维度 T * (D//H) = D，
                             因为这里 H=T 所以 mixed token 维度恰好还是 D)
        """
        B, T, D = x.shape
        assert T == self.T and D == self.D

        # Step1: SplitHead -> [B, T, H, head_dim]
        x = x.view(B, T, self.H, self.head_dim)

        # Step2: 交换 T 和 H 维度，实现"跨 token 按 head 重组"
        # 变换前 x[b, t, h, :] 表示第 t 个token的第 h 个head
        # 变换后 s[b, h, t, :] 把同一个 h 下所有 t 的 head 排在一起
        s = x.transpose(1, 2).contiguous()  # [B, H, T, head_dim]

        # Step3: 把 (T, head_dim) merge 成一个维度，得到每个 mixed token 的向量
        # s^h = Concat(x_1^h, x_2^h, ..., x_T^h)  ->  维度 T * head_dim
        s = s.view(B, self.H, T * self.head_dim)  # [B, H, T*head_dim]

        # 因为 H = T 且 T*head_dim = T*(D//T)，一般设计中会让 T*(D//T) = D
        # 若 D 恰好能被 T 整除，则 mixed token 维度与输入 D 完全一致
        return s  # [B, T(=H), D]

class PerTokenFFN(nn.Module):
    """
    对应论文 3.3.2 节 公式 (6)(7)(8)(9)。
    每个 token 拥有独立、不共享的两层 MLP 参数。
    用 batched matmul（一次性对所有 token 做不同的线性变换）实现，
    以保持GPU上的大GEMM并行效率。
    """

    def __init__(self, num_tokens: int, model_dim: int, k: int = 4):
        super().__init__()
        self.T = num_tokens
        self.D = model_dim
        self.hidden = k * model_dim

        # 用 [T, D, kD] 的参数张量表示"每个 token 独立的第一层权重"
        self.w1 = nn.Parameter(torch.empty(num_tokens, model_dim, self.hidden))
        self.b1 = nn.Parameter(torch.zeros(num_tokens, self.hidden))
        self.w2 = nn.Parameter(torch.empty(num_tokens, self.hidden, model_dim))
        self.b2 = nn.Parameter(torch.zeros(num_tokens, model_dim))

        nn.init.xavier_uniform_(self.w1)
        nn.init.xavier_uniform_(self.w2)

    def forward(self, s: torch.Tensor) -> torch.Tensor:
        """
        s: [B, T, D]
        return v: [B, T, D]
        """
        # 第一层：对每个 token t，做 s_t @ W1_t + b1_t
        # 用 einsum 一次性对所有 token 做 batched matmul，映射到工程上就是
        # 融合成一个大 kernel（"fusing parallel per-token FFNs into one kernel"），
        # 这正是论文提升 MFU 的关键工程手段之一。
        h = torch.einsum('btd,tdh->bth', s, self.w1) + self.b1  # [B, T, hidden]
        h = F.gelu(h)
        v = torch.einsum('bth,thd->btd', h, self.w2) + self.b2  # [B, T, D]
        return v

class Expert(nn.Module):
    """单个专家：标准两层MLP"""

    def __init__(self, model_dim: int, hidden_dim: int):
        super().__init__()
        self.fc1 = nn.Linear(model_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, model_dim)

    def forward(self, x):
        return self.fc2(F.gelu(self.fc1(x)))


class PerTokenSparseMoEFFN(nn.Module):
    """
    对应论文 3.4 节：ReLU Routing + DTSI-MoE (Dense-Training / Sparse-Inference)。

    关键点：
      1. 每个 token 有自己独立的一组专家（Per-token + MoE 组合）
      2. 路由用 ReLU(h(s_i)) 而不是 Top-k + Softmax，天然稀疏且可微
      3. 训练时使用两套路由器 h_train / h_infer，
         L_reg（稀疏正则）只作用于 h_infer，
         但两套路由器在训练阶段都会更新（Dense-Training）；
         推理阶段只使用 h_infer 做真正的稀疏计算（Sparse-Inference）
    """

    def __init__(self, num_tokens: int, model_dim: int, num_experts: int,
                 expert_hidden: int, reg_coef: float = 1e-3):
        super().__init__()
        self.T = num_tokens
        self.D = model_dim
        self.Ne = num_experts
        self.reg_coef = reg_coef

        # 每个 token 拥有自己独立的一组专家
        self.experts = nn.ModuleList([
            nn.ModuleList([Expert(model_dim, expert_hidden) for _ in range(num_experts)])
            for _ in range(num_tokens)
        ])

        # 两套路由器：训练用 / 推理用，每个 token 独立
        self.router_train = nn.Parameter(torch.empty(num_tokens, model_dim, num_experts))
        self.router_infer = nn.Parameter(torch.empty(num_tokens, model_dim, num_experts))
        nn.init.xavier_uniform_(self.router_train)
        nn.init.xavier_uniform_(self.router_infer)

        self.last_reg_loss = 0.0  # 供外部读取加到 total loss 里

    def _route(self, s: torch.Tensor, router: torch.Tensor) -> torch.Tensor:
        """
        s: [B, T, D], router: [T, D, Ne]
        return gate: [B, T, Ne]  (ReLU 之后, 天然带0，实现稀疏)
        对应公式(10): G_{i,j} = ReLU(h(s_i))
        """
        logits = torch.einsum('btd,tde->bte', s, router)
        return F.relu(logits)

    def forward(self, s: torch.Tensor) -> torch.Tensor:
        B, T, D = s.shape

        if self.training:
            # ---- Dense-Training：两套路由器都参与前向、都计算梯度 ----
            gate_train = self._route(s, self.router_train)   # [B,T,Ne]
            gate_infer = self._route(s, self.router_infer)   # [B,T,Ne]

            # 稀疏正则只施加在 infer 路由器上，对应公式(11)
            self.last_reg_loss = self.reg_coef * gate_infer.sum(dim=(1, 2)).mean()

            out_train = self._compute_experts_output(s, gate_train)
            out_infer = self._compute_experts_output(s, gate_infer)

            # 两路输出都参与训练（简单相加/平均均可，具体融合方式工程上可调），
            # 保证所有专家都获得充分梯度，避免"死亡专家"
            return 0.5 * (out_train + out_infer)
        else:
            # ---- Sparse-Inference：只用 infer 路由器，且可结合门控值做真稀疏跳过计算 ----
            gate_infer = self._route(s, self.router_infer)
            self.last_reg_loss = 0.0
            return self._compute_experts_output(s, gate_infer, sparse_exec=True)

    def _compute_experts_output(self, s: torch.Tensor, gate: torch.Tensor,
                                 sparse_exec: bool = False) -> torch.Tensor:
        """
        对应公式(10): v_i = sum_j G_{i,j} * e_{i,j}(s_i)
        sparse_exec=True 时，跳过 gate==0 的专家计算以节省真实推理算力
        （工程上通常按 batch 聚合非零 mask 做 gather/scatter，
         这里为了可读性用简单实现）
        """
        B, T, D = s.shape
        outputs = torch.zeros_like(s)
        for t in range(T):
            token_in = s[:, t, :]              # [B, D]
            token_gate = gate[:, t, :]          # [B, Ne]
            token_out = torch.zeros_like(token_in)
            for e_idx, expert in enumerate(self.experts[t]):
                g = token_gate[:, e_idx].unsqueeze(-1)  # [B, 1]
                if sparse_exec:
                    active_mask = (g.squeeze(-1) > 0)
                    if not active_mask.any():
                        continue  # 该专家在这个 batch 里对所有样本都未激活，跳过计算
                    active_idx = active_mask.nonzero(as_tuple=True)[0]
                    expert_out = expert(token_in[active_idx])
                    token_out[active_idx] += g[active_idx] * expert_out
                else:
                    token_out += g * expert(token_in)
            outputs[:, t, :] = token_out
        return outputs

class RankMixerBlock(nn.Module):
    """
    对应论文公式(1):
      S_{n-1} = LN(TokenMixing(X_{n-1}) + X_{n-1})
      X_n     = LN(PFFN(S_{n-1}) + S_{n-1})
    """

    def __init__(self, num_tokens: int, model_dim: int, k: int = 4,
                 use_moe: bool = False, num_experts: int = 4, expert_hidden: int = None):
        super().__init__()
        self.token_mixing = MultiHeadTokenMixing(num_tokens, model_dim)
        self.ln1 = nn.LayerNorm(model_dim)

        if use_moe:
            expert_hidden = expert_hidden or k * model_dim
            self.ffn = PerTokenSparseMoEFFN(num_tokens, model_dim, num_experts, expert_hidden)
        else:
            self.ffn = PerTokenFFN(num_tokens, model_dim, k)
        self.ln2 = nn.LayerNorm(model_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        s = self.ln1(self.token_mixing(x) + x)
        out = self.ln2(self.ffn(s) + s)
        return out


class RankMixer(nn.Module):
    """
    完整的 RankMixer 模型：Tokenization -> L层Block -> mean pooling -> 多任务头
    """

    def __init__(self, input_dim: int, num_tokens: int, model_dim: int,
                 num_layers: int, k: int = 4, use_moe: bool = False,
                 num_experts: int = 4, task_names=("finish", "skip", "like")):
        super().__init__()
        self.tokenizer = FeatureTokenizer(input_dim, num_tokens, model_dim)
        self.blocks = nn.ModuleList([
            RankMixerBlock(num_tokens, model_dim, k, use_moe, num_experts)
            for _ in range(num_layers)
        ])
        self.task_heads = nn.ModuleDict({
            name: nn.Linear(model_dim, 1) for name in task_names
        })

    def forward(self, e_input: torch.Tensor) -> dict:
        x = self.tokenizer(e_input)          # [B, T, D]
        for block in self.blocks:
            x = block(x)                     # [B, T, D]
        pooled = x.mean(dim=1)               # mean pooling -> [B, D]
        return {name: torch.sigmoid(head(pooled)).squeeze(-1)
                for name, head in self.task_heads.items()}


# ------------------- 简单自测 -------------------
if __name__ == "__main__":
    B, T, D, L = 8, 16, 768, 2   # 对应论文100M配置 D=768,T=16,L=2
    input_dim = T * 64           # 假设每个token切分前维度d=64

    model = RankMixer(input_dim=input_dim, num_tokens=T, model_dim=D,
                       num_layers=L, k=4, use_moe=False)
    dummy_input = torch.randn(B, input_dim)
    outputs = model(dummy_input)
    for task, pred in outputs.items():
        print(task, pred.shape)  # 期望输出: torch.Size([8])