     for (int i = threadIdx.x; i < lastValid; i += TPB)
     {
         const int idx = offset + i;
-        threadData = input[idx];
+        threadData = max(static_cast<float>(input[idx]), threadData);
     }
 
     const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, cub::Max());

    for (int i = threadIdx.x; i < ld; i += TPB)
    {
        const int idx = offset + i;
        const float val = (i < lastValid) ? exp(float(input[idx]) * w) * rZ : 0.f;
        output[idx] = T(val);
    }
}

    }
    __syncthreads();

    threadData = 0;

    for (int i = threadIdx.x; i < lastValid; i += TPB)
    {
        const int idx = offset + i;
        threadData += exp((static_cast<float>(input[idx]) - fMax) * w);
    }

    const auto Z = BlockReduce(tmpStorage).Reduce(threadData, sum);
