// Fused ensemble reduce kernel: sigmoid + per-model mean + weighted sum. // // One thread per model. Each thread computes sigmoid(logit) for its segment, // averages, multiplies by the normalized weight, and atomically adds to a // single output scalar. The output is the weighted mean of sigmoid activations // across all models -- a value in [0, 1] that the host remaps to [-1, 1]. // // All data is f32 (logits come from CPU-side Vec, not bf16 GpuTensor). extern "C" __global__ void ensemble_sigmoid_mean_reduce( const float* __restrict__ logits, // [total_logits] concatenated const int* __restrict__ offsets, // [n_models] start offset per model const int* __restrict__ lengths, // [n_models] logit count per model const float* __restrict__ weights, // [n_models] normalized weights float* __restrict__ output, // [1] result scalar (must be zeroed) int n_models) { int model = blockIdx.x * blockDim.x + threadIdx.x; if (model >= n_models) return; int off = offsets[model]; int len = lengths[model]; if (len <= 0) return; // Fused sigmoid + accumulate for this model's segment float sum = 0.0f; for (int j = 0; j < len; ++j) { float x = logits[off + j]; // sigmoid(x) = 1 / (1 + exp(-x)) float s = 1.0f / (1.0f + expf(-x)); sum += s; } float mean_sigmoid = sum / (float)len; // Weighted contribution (weights are pre-normalized on host) float contribution = mean_sigmoid * weights[model]; atomicAdd(output, contribution); }