perf: RepEncoder 批量 embedding 查表（28 次 kernel launch → 1 次）

所有 slot 的 sign id 合并为一次 embedding lookup，再按 slot 拆分做 segment_reduce。
数学等价，纯 GPU 算子优化。

代码/code/infer.py

@@ -257,16 +257,33 @@ class RepEncoder(nn.Module):
         self.linear = nn.Linear(in_features=slot_num * emb_dim, out_features=d_model)
 
     def forward(self, batch):
-        pooled_embs = []
         max_idx = self.emb.num_embeddings - 1
         target_dtype = self.input_norm.weight.dtype  # 后续层 dtype（FP16 时为 torch.float16）
+
+        # 批量收集所有 slot 的 values，一次 embedding 查表（减少 28 → 1 次 kernel launch）
+        all_values = []
+        all_offsets = []
+        slot_boundaries = [0]  # 记录每个 slot 在 all_values 中的起止位置
         for i in range(self.slot_num):
             values, offsets = batch[i + 1]
             offsets = offsets.to(values.device)
-            values = values.clamp(0, max_idx)  # 超出 vocab_size 的 sign id 截断，避免越界
-            sign_emb = self.emb(values).to(target_dtype)
-            res = torch.segment_reduce(sign_emb, reduce='sum', offsets=offsets, initial=0)
+            values = values.clamp(0, max_idx)
+            all_values.append(values)
+            all_offsets.append(offsets)
+            slot_boundaries.append(slot_boundaries[-1] + values.size(0))
+
+        # 一次批量 embedding 查表
+        values_cat = torch.cat(all_values)
+        embs_cat = self.emb(values_cat).to(target_dtype)
+
+        # 按 slot 拆分并 segment_reduce
+        pooled_embs = []
+        for i in range(self.slot_num):
+            start, end = slot_boundaries[i], slot_boundaries[i + 1]
+            slot_embs = embs_cat[start:end]
+            res = torch.segment_reduce(slot_embs, reduce='sum', offsets=all_offsets[i], initial=0)
             pooled_embs.append(res)
+
         fused_embs = torch.cat(pooled_embs, dim=1)
         norm_emb = self.input_norm(fused_embs)
         rep_emb = self.linear(norm_emb)