2.2 Reduction、Scan 与 Softmax 的共同骨架

每线程处理若干元素 → Block 内树形合并 → 跨 Block 二次归约
Scan 多一个“前缀传播”阶段；Softmax 用 max 与 sum 两次归约

import torch

def stable_softmax(x, dim=-1):
    m = x.max(dim=dim, keepdim=True).values
    e = torch.exp((x - m).float())
    return e / e.sum(dim=dim, keepdim=True)

def inclusive_scan(x):
    return torch.cumsum(x, dim=-1)

shared[threadIdx.x] = local_sum;
__syncthreads();
for (int offset = blockDim.x / 2; offset > 0; offset /= 2) {
    if (threadIdx.x < offset) shared[threadIdx.x] += shared[threadIdx.x + offset];
    __syncthreads();
}
if (threadIdx.x == 0) partial[blockIdx.x] = shared[0];