perf: searchsorted 构造因果mask，消除最后一个同步点(repeat_interleave张量repeats)

dense MoE 去掉MoE的nonzero同步省了评测20s；embedding融合(无同步)只省1s
->真正的杠杆是消同步点。mask构造的repeat_interleave(lengths张量)是model(batch)
内最后一个同步点，改用searchsorted求doc_id(输出size已知,无同步)。等价测试已加。

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

代码/code/infer.py

@@ -49,6 +49,7 @@ CONFIG = {
     # 须靠提交验证。AUC中性、MoE仅占2%算力故风险极低。
     "vectorize_moe": True,    # True=稠密向量化MoE(无同步点)；False=原逐expert循环(.nonzero同步)
     "fuse_embedding": True,   # True=28个slot的查表+池化融合为1次(减per-batch kernel启动)
+    "syncfree_mask": True,    # True=用searchsorted构造因果mask(无同步)；False=repeat_interleave(同步)
     "compile": False,         # 是否 torch.compile（实测慢5×，勿开）
 }
 
@@ -596,11 +597,20 @@ class CTRModel(nn.Module):
         lengths = seq_info[1:] - seq_info[:-1]
         lengths = lengths.view(-1)
         indices = torch.cumsum(torch.ones_like(lengths), dim=0) - 1
-        result = torch.repeat_interleave(indices, lengths)
+        result = torch.repeat_interleave(indices, lengths)  # repeats 是张量 → 同步
         a = result.view(1, -1) - result.view(-1, 1)
         out_mask = torch.tril((a == 0).to(torch.int32)).bool()
         return out_mask
 
+    def causal_mask_syncfree(self, user_offsets, S, device):
+        """与 get_sequence_causal_mask 等价，但用 searchsorted 求每个位置的用户号，
+        避免 repeat_interleave(张量repeats) 的隐式同步。"""
+        pos = torch.arange(S, device=device)
+        doc_id = torch.searchsorted(user_offsets[1:].contiguous(), pos, right=True)  # [S]，无同步
+        same = doc_id.view(-1, 1) == doc_id.view(1, -1)
+        causal = pos.view(-1, 1) >= pos.view(1, -1)
+        return same & causal
+
     def build_block_mask(self, user_offsets, S):
         """FlexAttention 块对角因果 mask：q 只能 attend 同一用户且 kv<=q 的位置。"""
         lengths = (user_offsets[1:] - user_offsets[:-1]).view(-1)
@@ -623,7 +633,11 @@ class CTRModel(nn.Module):
             S = seq_input.shape[0]  # rep_encoder 输出 [S, D]，S=总 token 数
             extension = {"block_mask": self.build_block_mask(user_offsets, S)}
         else:
-            seq_mask = self.get_sequence_causal_mask(user_offsets)
+            if CONFIG.get("syncfree_mask", True):
+                seq_mask = self.causal_mask_syncfree(
+                    user_offsets, seq_input.shape[0], seq_input.device)
+            else:
+                seq_mask = self.get_sequence_causal_mask(user_offsets)
             extension = {"mask": seq_mask.unsqueeze(0).unsqueeze(0)}
         encoder_output, moe_loss = self.seq_encoder(x=seq_input, extension=extension)
         encoder_output = encoder_output.squeeze(0)