Adder Transformer Inference

原始题目：LeetGPU - Adder Transformer Inference

题目描述

对一个仅有 10 个参数的微型 Transformer 执行批量自回归推理，该模型能够以 $\ge 99\%$ 的准确率相加两个 10 位十进制数。给定形状为 $[batch\_size, 31]$ 的 int32 提示和包含 10 个 float32 权重的缓冲区，输出形状为 $[batch\_size, 11, 10]$ 的 logits——每个解码步对应一行，覆盖 10 个数字类别（0–9）。

模型来自 AdderBoard 竞赛，通过 RoPE 几何编码、tied embeddings 和 SwiGLU 门控在 10 个参数中编码进位传播。

架构：单层 pre-norm Transformer，hidden_dim=2，1 head，head_dim=2，词表=10。每步对序列 $[B, seq\_len, 2]$ 运行完整前向传播：Token Embedding → RMSNorm → Self-Attention（Q/K/V 投影 + RoPE + Causal Attn + 残差） → RMSNorm → MLP（Gate + SwiGLU + Carry + 残差） → Final RMSNorm → 输出 Logits。

从 31-token 提示开始，重复 11 次解码步，序列长度从 31 增长到 42。

示例

Input:  batch_size=2, pairs (3+5), (99+1)
        prompts[0] = [0,3,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,5,0,0,0,0,0,0,0,0,0,0]
        prompts[1] = [0,9,9,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0]
Output: logits shape [2, 11, 10]
        Pair (3,5):  sum=8    → argmax: [8,0,0,0,0,0,0,0,0,0,0]
        Pair (99,1): sum=100  → argmax: [0,0,1,0,0,0,0,0,0,0,0]

约束条件

• $1 \le batch\_size \le 100{,}000$
• prompts: int32，值 ∈ [0, 9]；weights: float32，恰好 10 个元素
• 输入数字范围 $[0, 9{,}999{,}999{,}999]$（10 位无符号整数）
• 架构常量固定：vocab_size=10, hidden_dim=2, head_dim=2, num_heads=1, prompt_len=31, decode_steps=11
• RMSNorm $\epsilon = 10^{-6}$，RoPE $\omega = 2\pi/19$
• 性能测试在 $batch\_size = 100{,}000$ 下进行

解题思路

这道题考验的是在 GPU 上高效实现完整 Transformer 推理 pipeline 的能力。虽然模型极小（10 个参数），但 $batch\_size$ 可达 10 万，需要高度并行化。关键优化包括：将权重预加载到常量/共享内存、融合逐元素操作（RMSNorm、SiLU）、减少中间张量的内存分配。11 步自回归循环的串行依赖是性能瓶颈——每步都依赖上一步的 argmax 输出。欢迎在 GitHub Discussions 分享你的解法。