Softmax · LLM 算子刷题

给定一个含 N 个 float 的数组 x，把它的 softmax 写进 y：

y[i] = exp(x[i]) / (exp(x[0]) + exp(x[1]) + ... + exp(x[N-1]))

每个 y[i] 都在 [0, 1] 之间，全部加起来等于 1。float32 下 exp 的输入超过约 88.7 就会溢出成无穷大，无穷大除以无穷大得 NaN。每个 x[i] 减去同一个常数，y[i] 不变，所以先减去 x 的最大值 m：

y[i] = exp(x[i] - m) / (exp(x[0] - m) + ... + exp(x[N-1] - m))

这样每个指数都不超过 0，不会溢出。

要求

在所选框架的初始代码里实现 solve，签名保持不变。
- CUDA：extern "C" void solve(const float* x, float* y, int N)
- Triton、PyTorch：def solve(x: torch.Tensor, y: torch.Tensor, N: int)
x、y 已经在 GPU 上。不要修改 x。
结果原地写进 y；进入 solve 时 y 里的内容是未定义的。
可以在 solve 里申请临时显存。

输入：x = [0.0, 1.0, 2.0, 3.0], N = 4
输出：y = [0.032059, 0.087144, 0.236883, 0.643914]

输入：x = [100.0, 100.0, 99.0], N = 3
输出：y = [0.422319, 0.422319, 0.155362]

第二个示例里，float32 的 exp(100.0) 已经是无穷大；减去 m = 100.0 后，指数是 0、0、-1。