#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <math_constants.h>
#include <stdint.h>

struct __align__(4) block_q8_1 {
    half d;
    half s;
    int8_t qs[32];
};

static_assert(sizeof(block_q8_1) == 36, "wrong q8_1 block size");

__device__ __forceinline__ float makepad_cuda_bf16_to_f32(uint16_t word) {
    const uint32_t bits = static_cast<uint32_t>(word) << 16;
    return __uint_as_float(bits);
}

__device__ __forceinline__ float makepad_cuda_f16_to_f32(uint16_t word) {
    return __half2float(*reinterpret_cast<const half *>(&word));
}

__device__ __forceinline__ uint16_t makepad_cuda_f32_to_f16_bits(float value) {
    const half h = __float2half_rn(value);
    return *reinterpret_cast<const uint16_t *>(&h);
}

__device__ __forceinline__ uint16_t makepad_cuda_f32_to_bf16_bits(float value) {
    const uint32_t bits = __float_as_uint(value);
    return static_cast<uint16_t>(bits >> 16);
}

__device__ __forceinline__ float makepad_cuda_bf16_round(float value) {
    const uint32_t bits = __float_as_uint(value);
    return __uint_as_float(bits & 0xFFFF0000u);
}

template <typename T>
__device__ __forceinline__ T makepad_cuda_warp_reduce_sum(T value) {
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        value += __shfl_down_sync(0xffffffffu, value, offset);
    }
    return value;
}

template <typename T>
__device__ __forceinline__ T makepad_cuda_warp_reduce_max(T value) {
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        const T other = __shfl_down_sync(0xffffffffu, value, offset);
        value = value > other ? value : other;
    }
    return value;
}

template <typename T>
__device__ __forceinline__ T makepad_cuda_block_reduce_sum(T value) {
    __shared__ T shared[32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    value = makepad_cuda_warp_reduce_sum(value);
    if (lane == 0) {
        shared[warp] = value;
    }
    __syncthreads();
    value = threadIdx.x < (blockDim.x + 31) / 32 ? shared[lane] : T(0);
    if (warp == 0) {
        value = makepad_cuda_warp_reduce_sum(value);
    }
    return value;
}

template <typename T>
__device__ __forceinline__ T makepad_cuda_block_reduce_max(T value) {
    __shared__ T shared[32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    value = makepad_cuda_warp_reduce_max(value);
    if (lane == 0) {
        shared[warp] = value;
    }
    __syncthreads();
    value = threadIdx.x < (blockDim.x + 31) / 32 ? shared[lane] : -CUDART_INF_F;
    if (warp == 0) {
        value = makepad_cuda_warp_reduce_max(value);
    }
    return value;
}

static __global__ void makepad_cuda_quantize_q8_1_f32_kernel(
        const float * __restrict__ input,
        block_q8_1 * __restrict__ output,
        uint32_t block_count) {
    const uint32_t block_idx = blockIdx.x;
    const uint32_t lane = threadIdx.x;
    if (block_idx >= block_count || lane >= 32) {
        return;
    }

    const float xi = input[block_idx * 32 + lane];
    float amax = fabsf(xi);
    amax = makepad_cuda_warp_reduce_max(amax);
    const float d = amax / 127.0f;
    const float id = d != 0.0f ? 1.0f / d : 0.0f;
    const int8_t q = amax == 0.0f ? 0 : static_cast<int8_t>(lrintf(xi * id));
    int sum = static_cast<int>(q);
    sum = makepad_cuda_warp_reduce_sum(sum);
    output[block_idx].qs[lane] = q;
    if (lane == 0) {
        output[block_idx].d = __float2half_rn(d);
        output[block_idx].s = __float2half_rn(static_cast<float>(sum) * d);
    }
}

static __global__ void makepad_cuda_scale_f32_kernel(
        float * __restrict__ values,
        float scale,
        uint32_t n) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) {
        return;
    }
    values[idx] = makepad_cuda_bf16_round(values[idx] * scale);
}

static __global__ void makepad_cuda_scale_f32_device_f32_index_kernel(
        float * __restrict__ values,
        const float * __restrict__ scales,
        uint32_t scale_index,
        uint32_t n) {
    const float scale = scales[scale_index];
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) {
        return;
    }
    values[idx] = makepad_cuda_bf16_round(values[idx] * scale);
}

static __global__ void makepad_cuda_f32_to_bf16_kernel(
        const float * __restrict__ input,
        uint16_t * __restrict__ output,
        uint32_t n) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) {
        return;
    }
    output[idx] = makepad_cuda_f32_to_bf16_bits(input[idx]);
}

static __global__ void makepad_cuda_bf16_to_f32_kernel(
        const uint16_t * __restrict__ input,
        float * __restrict__ output,
        uint32_t n) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) {
        return;
    }
    output[idx] = makepad_cuda_bf16_to_f32(input[idx]);
}

// Exact layout of Michelangelo FourierEmbedder(num_freqs=8,
// include_input=true, include_pi=false) followed by concatenated normals:
// [xyz, sin(xyz * 2^[0..7]), cos(xyz * 2^[0..7]), normals].
static __global__ void makepad_cuda_skintokens_michelangelo_fourier_f32_kernel(
        const float * __restrict__ condition,
        float * __restrict__ output,
        uint32_t rows) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    const uint32_t values = rows * 54u;
    if (idx >= values) {
        return;
    }
    const uint32_t row = idx / 54u;
    const uint32_t col = idx - row * 54u;
    const float *input_row = condition + row * 6u;
    if (col < 3u) {
        output[idx] = input_row[col];
        return;
    }
    if (col >= 51u) {
        output[idx] = input_row[col - 48u];
        return;
    }
    const uint32_t trig_col = col - 3u;
    const bool cosine = trig_col >= 24u;
    const uint32_t component = trig_col % 24u;
    const uint32_t coordinate = component / 8u;
    const uint32_t frequency = component % 8u;
    const float angle = __fmul_rn(input_row[coordinate], ldexpf(1.0f, frequency));
    output[idx] = cosine ? cosf(angle) : sinf(angle);
}

static __global__ void makepad_cuda_add_f32_kernel(
        const float * __restrict__ left,
        const float * __restrict__ right,
        float * __restrict__ out,
        uint32_t n) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) {
        return;
    }
    out[idx] = makepad_cuda_bf16_round(left[idx] + right[idx]);
}

static __global__ void makepad_cuda_copy_f32_kernel(
        const float * __restrict__ input,
        float * __restrict__ output,
        uint32_t n) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) {
        return;
    }
    output[idx] = input[idx];
}

static __global__ void makepad_cuda_weighted_sum_rows_f32_kernel(
        const float * __restrict__ batched_inputs,
        const float * __restrict__ weights,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t input_count) {
    const uint32_t row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= row_count) {
        return;
    }
    float total = 0.0f;
    for (uint32_t slot = 0; slot < input_count; ++slot) {
        total += batched_inputs[slot * row_count + row] * weights[slot];
    }
    output[row] = makepad_cuda_bf16_round(total);
}

static __global__ void makepad_cuda_weighted_sum_rows_grouped_f32_kernel(
        const float * __restrict__ batched_inputs,
        const float * __restrict__ weights,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t input_count) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    const uint32_t total = row_count * row_stride;
    if (idx >= total) {
        return;
    }
    const uint32_t row = idx / row_stride;
    const uint32_t col = idx - row * row_stride;
    float accum = 0.0f;
    for (uint32_t slot = 0; slot < input_count; ++slot) {
        const uint32_t input_idx = slot * total + row * row_stride + col;
        accum += batched_inputs[input_idx] * weights[row * input_count + slot];
    }
    output[idx] = makepad_cuda_bf16_round(accum);
}

static __global__ void makepad_cuda_add_scaled_rows_f32_kernel(
        const float * __restrict__ input,
        const float * __restrict__ scales,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t row_stride) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    const uint32_t total = row_count * row_stride;
    if (idx >= total) {
        return;
    }
    const uint32_t row = idx / row_stride;
    const float product = makepad_cuda_bf16_round(input[idx] * scales[row]);
    output[idx] = makepad_cuda_bf16_round(output[idx] + product);
}

static __global__ void makepad_cuda_add_scaled_rows_f32_indexed_kernel(
        const float * __restrict__ input,
        const float * __restrict__ scales,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t scale_row_stride,
        uint32_t scale_column) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    const uint32_t total = row_count * row_stride;
    if (idx >= total) {
        return;
    }
    const uint32_t row = idx / row_stride;
    const float scale = scales[row * scale_row_stride + scale_column];
    const float product = makepad_cuda_bf16_round(input[idx] * scale);
    output[idx] = makepad_cuda_bf16_round(output[idx] + product);
}

static __global__ void makepad_cuda_mul_f32_kernel(
        const float * __restrict__ left,
        const float * __restrict__ right,
        float * __restrict__ out,
        uint32_t n) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) {
        return;
    }
    out[idx] = makepad_cuda_bf16_round(left[idx] * right[idx]);
}

static __global__ void makepad_cuda_gelu_f32_kernel(
        const float * __restrict__ input,
        float * __restrict__ out,
        uint32_t n) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) {
        return;
    }
    const float value = input[idx];
    const float squared = makepad_cuda_bf16_round(value * value);
    const float cubic = makepad_cuda_bf16_round(squared * value);
    const float poly = makepad_cuda_bf16_round(value + makepad_cuda_bf16_round(0.044715f * cubic));
    const float tanh_input = makepad_cuda_bf16_round(0.7978846f * poly);
    const float tanh_value = makepad_cuda_bf16_round(tanhf(tanh_input));
    const float half = makepad_cuda_bf16_round(0.5f * value);
    out[idx] = makepad_cuda_bf16_round(half * makepad_cuda_bf16_round(1.0f + tanh_value));
}

static __global__ void makepad_cuda_geglu_split_f32_kernel(
        const float * __restrict__ gate_up,
        float * __restrict__ out,
        uint32_t n,
        uint32_t split_offset) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) {
        return;
    }
    const float gate = gate_up[idx];
    const float up = gate_up[split_offset + idx];
    const float squared = makepad_cuda_bf16_round(gate * gate);
    const float cubic = makepad_cuda_bf16_round(squared * gate);
    const float poly = makepad_cuda_bf16_round(gate + makepad_cuda_bf16_round(0.044715f * cubic));
    const float tanh_input = makepad_cuda_bf16_round(0.7978846f * poly);
    const float tanh_value = makepad_cuda_bf16_round(tanhf(tanh_input));
    const float half = makepad_cuda_bf16_round(0.5f * gate);
    const float gelu = makepad_cuda_bf16_round(half * makepad_cuda_bf16_round(1.0f + tanh_value));
    out[idx] = makepad_cuda_bf16_round(gelu * up);
}

static __global__ void makepad_cuda_geglu_split_f32_rows_kernel(
        const float * __restrict__ gate_up,
        float * __restrict__ out,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        uint32_t split_offset) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    const uint32_t total = row_count * n;
    if (idx >= total) {
        return;
    }
    const uint32_t row = idx / n;
    const uint32_t col = idx % n;
    const float * row_in = gate_up + row * row_stride;
    const float gate = row_in[col];
    const float up = row_in[split_offset + col];
    const float squared = makepad_cuda_bf16_round(gate * gate);
    const float cubic = makepad_cuda_bf16_round(squared * gate);
    const float poly = makepad_cuda_bf16_round(gate + makepad_cuda_bf16_round(0.044715f * cubic));
    const float tanh_input = makepad_cuda_bf16_round(0.7978846f * poly);
    const float tanh_value = makepad_cuda_bf16_round(tanhf(tanh_input));
    const float half = makepad_cuda_bf16_round(0.5f * gate);
    const float gelu = makepad_cuda_bf16_round(half * makepad_cuda_bf16_round(1.0f + tanh_value));
    out[idx] = makepad_cuda_bf16_round(gelu * up);
}

static __global__ void makepad_cuda_rms_norm_row_weighted_f32_kernel(
        const float * __restrict__ input,
        const uint16_t * __restrict__ weights_bf16,
        float * __restrict__ output,
        uint32_t n,
        float eps) {
    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float v = input[idx];
        sum += v * v;
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(n) + eps);
    }
    __syncthreads();
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float normalized = makepad_cuda_bf16_round(input[idx] * inv_rms);
        const float weight = makepad_cuda_bf16_to_f32(weights_bf16[idx]);
        output[idx] = makepad_cuda_bf16_round(normalized * weight);
    }
}

static __global__ void makepad_cuda_rms_norm_row_weighted_f32_f32weights_kernel(
        const float * __restrict__ input,
        const float * __restrict__ weights_f32,
        float * __restrict__ output,
        uint32_t n,
        float eps) {
    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float v = input[idx];
        sum += v * v;
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(n) + eps);
    }
    __syncthreads();
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float normalized = makepad_cuda_bf16_round(input[idx] * inv_rms);
        output[idx] = makepad_cuda_bf16_round(normalized * weights_f32[idx]);
    }
}

static __global__ void makepad_cuda_rms_norm_row_weighted_f32_f32weights_precise_kernel(
        const float * __restrict__ input,
        const float * __restrict__ weights_f32,
        float * __restrict__ output,
        uint32_t n,
        float eps) {
    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float v = input[idx];
        sum += v * v;
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(n) + eps);
    }
    __syncthreads();
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        output[idx] = input[idx] * inv_rms * weights_f32[idx];
    }
}

static __global__ void makepad_cuda_rms_norm_rows_weighted_f32_kernel(
        const float * __restrict__ input,
        const uint16_t * __restrict__ weights_bf16,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps) {
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const float * row_in = input + row * row_stride;
    float * row_out = output + row * row_stride;
    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float v = row_in[idx];
        sum += v * v;
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(n) + eps);
    }
    __syncthreads();
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float normalized = makepad_cuda_bf16_round(row_in[idx] * inv_rms);
        const float weight = makepad_cuda_bf16_to_f32(weights_bf16[idx]);
        row_out[idx] = makepad_cuda_bf16_round(normalized * weight);
    }
}

static __global__ void makepad_cuda_rms_norm_rows_weighted_f32_f32weights_kernel(
        const float * __restrict__ input,
        const float * __restrict__ weights_f32,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps) {
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const float * row_in = input + row * row_stride;
    float * row_out = output + row * row_stride;
    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float v = row_in[idx];
        sum += v * v;
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(n) + eps);
    }
    __syncthreads();
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float normalized = makepad_cuda_bf16_round(row_in[idx] * inv_rms);
        row_out[idx] = makepad_cuda_bf16_round(normalized * weights_f32[idx]);
    }
}

static __global__ void makepad_cuda_rms_norm_rows_weighted_f32_f32weights_precise_kernel(
        const float * __restrict__ input,
        const float * __restrict__ weights_f32,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps) {
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const float * row_in = input + row * row_stride;
    float * row_out = output + row * row_stride;
    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float v = row_in[idx];
        sum += v * v;
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(n) + eps);
    }
    __syncthreads();
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        row_out[idx] = row_in[idx] * inv_rms * weights_f32[idx];
    }
}

// The precise weighted RMS-norm above, reading its groups out of a bf16
// column slab (flux2's fused qkv linear output kept in bf16 storage) and
// writing contiguous f32. Values expand losslessly, the strided per-thread
// accumulation and block reduction are identical, so the output bits match
// slice+expand+rms exactly. `groups_per_row` groups of `n` sit at
// `col_off` within each `in_stride`-wide row.
static __global__ void makepad_cuda_rms_norm_weighted_bf16slab_f32_kernel(
        const uint16_t * __restrict__ input,
        const float * __restrict__ weights_f32,
        float * __restrict__ output,
        uint32_t group_count,
        uint32_t groups_per_row,
        uint32_t in_stride,
        uint32_t col_off,
        uint32_t n,
        float eps) {
    const uint32_t group = blockIdx.x;
    if (group >= group_count) {
        return;
    }
    const uint32_t row = group / groups_per_row;
    const uint32_t sub = group - row * groups_per_row;
    const uint16_t * group_in =
        input + static_cast<size_t>(row) * in_stride + col_off + static_cast<size_t>(sub) * n;
    float * group_out = output + static_cast<size_t>(group) * n;
    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float v = makepad_cuda_bf16_to_f32(group_in[idx]);
        sum += v * v;
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(n) + eps);
    }
    __syncthreads();
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        group_out[idx] =
            makepad_cuda_bf16_to_f32(group_in[idx]) * inv_rms * weights_f32[idx];
    }
}

extern "C" cudaError_t makepad_cuda_rms_norm_weighted_bf16slab_f32(
        const uint16_t * input,
        const float * weights_f32,
        float * output,
        uint32_t group_count,
        uint32_t groups_per_row,
        uint32_t in_stride,
        uint32_t col_off,
        uint32_t n,
        float eps,
        cudaStream_t stream) {
    if (group_count == 0 || groups_per_row == 0 || n == 0
        || col_off + groups_per_row * n > in_stride) {
        return cudaErrorInvalidValue;
    }
    const uint32_t block = n < 1024 ? 256 : 1024;
    makepad_cuda_rms_norm_weighted_bf16slab_f32_kernel<<<group_count, block, 0, stream>>>(
        input, weights_f32, output, group_count, groups_per_row, in_stride, col_off, n, eps);
    return cudaGetLastError();
}

// Official Qwen3RMSNorm: x.to(f32); xhat=x*rsqrt; return w * xhat.to(input_dtype).
// Cast xhat to bf16, then multiply by w in f32. Do not bf16-round the product
// (that extra round is gpu_rms_norm_mul_bf16 and does not match the hook dump).
static __global__ void makepad_cuda_rms_norm_qwen3_kernel(
        const float * __restrict__ input,
        const float * __restrict__ weights_f32,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps) {
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const float * row_in = input + row * row_stride;
    float * row_out = output + row * row_stride;
    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float v = row_in[idx];
        sum += v * v;
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(n) + eps);
    }
    __syncthreads();
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float xhat = makepad_cuda_bf16_round(row_in[idx] * inv_rms);
        row_out[idx] = weights_f32[idx] * xhat;
    }
}

static __global__ void makepad_cuda_rms_norm_rows_no_scale_f32_kernel(
        const float * __restrict__ input,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps) {
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const float * row_in = input + row * row_stride;
    float * row_out = output + row * row_stride;
    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float v = row_in[idx];
        sum += v * v;
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(n) + eps);
    }
    __syncthreads();
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        row_out[idx] = makepad_cuda_bf16_round(row_in[idx] * inv_rms);
    }
}

static __global__ void makepad_cuda_rms_norm_rows_no_scale_f32_precise_kernel(
        const float * __restrict__ input,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps) {
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const float * row_in = input + row * row_stride;
    float * row_out = output + row * row_stride;
    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float v = row_in[idx];
        sum += v * v;
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(n) + eps);
    }
    __syncthreads();
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        row_out[idx] = row_in[idx] * inv_rms;
    }
}

static __global__ void makepad_cuda_rope_rows_f32_kernel(
        const float * __restrict__ input,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t head_dim,
        uint32_t rotary_dim,
        float base,
        uint32_t position) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    const uint32_t total = row_count * head_dim;
    if (idx >= total) {
        return;
    }
    const uint32_t row = idx / row_stride;
    const uint32_t col = idx % row_stride;
    if (row >= row_count || col >= head_dim) {
        return;
    }
    const uint32_t half = head_dim / 2;
    const uint32_t rotary_pairs = rotary_dim / 2;
    const float * row_in = input + row * row_stride;
    float * row_out = output + row * row_stride;

    if (col < rotary_pairs) {
        const float exponent = (2.0f * static_cast<float>(col)) / static_cast<float>(head_dim);
        const float inv_freq = powf(base, -exponent);
        const float theta = static_cast<float>(position) * inv_freq;
        const float cos_theta = cosf(theta);
        const float sin_theta = sinf(theta);
        const float left = row_in[col];
        const float right = row_in[half + col];
        row_out[col] = makepad_cuda_bf16_round(left * cos_theta - right * sin_theta);
    } else if (col >= half && col < half + rotary_pairs) {
        const uint32_t pair = col - half;
        const float exponent = (2.0f * static_cast<float>(pair)) / static_cast<float>(head_dim);
        const float inv_freq = powf(base, -exponent);
        const float theta = static_cast<float>(position) * inv_freq;
        const float cos_theta = cosf(theta);
        const float sin_theta = sinf(theta);
        const float left = row_in[pair];
        const float right = row_in[col];
        row_out[col] = makepad_cuda_bf16_round(left * sin_theta + right * cos_theta);
    } else {
        row_out[col] = row_in[col];
    }
}

static __global__ void makepad_cuda_rope_rows_f32_device_u32_kernel(
        const float * __restrict__ input,
        float * __restrict__ output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t head_dim,
        uint32_t rotary_dim,
        float base,
        const uint32_t * __restrict__ position_device_u32) {
    const uint32_t position = *position_device_u32;
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    const uint32_t total = row_count * head_dim;
    if (idx >= total) {
        return;
    }
    const uint32_t row = idx / row_stride;
    const uint32_t col = idx % row_stride;
    if (row >= row_count || col >= head_dim) {
        return;
    }
    const uint32_t half = head_dim / 2;
    const uint32_t rotary_pairs = rotary_dim / 2;
    const float * row_in = input + row * row_stride;
    float * row_out = output + row * row_stride;

    if (col < rotary_pairs) {
        const float exponent = (2.0f * static_cast<float>(col)) / static_cast<float>(head_dim);
        const float inv_freq = powf(base, -exponent);
        const float theta = static_cast<float>(position) * inv_freq;
        const float cos_theta = cosf(theta);
        const float sin_theta = sinf(theta);
        const float left = row_in[col];
        const float right = row_in[half + col];
        row_out[col] = makepad_cuda_bf16_round(left * cos_theta - right * sin_theta);
    } else if (col >= half && col < half + rotary_pairs) {
        const uint32_t pair = col - half;
        const float exponent = (2.0f * static_cast<float>(pair)) / static_cast<float>(head_dim);
        const float inv_freq = powf(base, -exponent);
        const float theta = static_cast<float>(position) * inv_freq;
        const float cos_theta = cosf(theta);
        const float sin_theta = sinf(theta);
        const float left = row_in[pair];
        const float right = row_in[col];
        row_out[col] = makepad_cuda_bf16_round(left * sin_theta + right * cos_theta);
    } else {
        row_out[col] = row_in[col];
    }
}

static __global__ void makepad_cuda_kv_append_f32_kernel(
        const float * __restrict__ keys,
        const float * __restrict__ values,
        uint16_t * __restrict__ key_cache,
        uint16_t * __restrict__ value_cache,
        uint32_t kv_head_count,
        uint32_t head_dim,
        uint32_t max_tokens,
        uint32_t slot) {
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    const uint32_t total = kv_head_count * head_dim;
    if (idx >= total) {
        return;
    }
    const uint32_t head = idx / head_dim;
    const uint32_t dim = idx % head_dim;
    const uint32_t row_base = head * max_tokens * head_dim;
    key_cache[row_base + slot * head_dim + dim] = makepad_cuda_f32_to_bf16_bits(keys[idx]);
    value_cache[row_base + dim * max_tokens + slot] = makepad_cuda_f32_to_bf16_bits(values[idx]);
}

static __global__ void makepad_cuda_kv_append_f32_device_u32_kernel(
        const float * __restrict__ keys,
        const float * __restrict__ values,
        uint16_t * __restrict__ key_cache,
        uint16_t * __restrict__ value_cache,
        uint32_t kv_head_count,
        uint32_t head_dim,
        uint32_t max_tokens,
        const uint32_t * __restrict__ slot_device_u32) {
    const uint32_t slot = *slot_device_u32;
    if (slot >= max_tokens) {
        return;
    }
    const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    const uint32_t total = kv_head_count * head_dim;
    if (idx >= total) {
        return;
    }
    const uint32_t head = idx / head_dim;
    const uint32_t dim = idx % head_dim;
    const uint32_t row_base = head * max_tokens * head_dim;
    key_cache[row_base + slot * head_dim + dim] = makepad_cuda_f32_to_bf16_bits(keys[idx]);
    value_cache[row_base + dim * max_tokens + slot] = makepad_cuda_f32_to_bf16_bits(values[idx]);
}

static __device__ __forceinline__ float makepad_cuda_weighted_norm_f32(
        const float * __restrict__ row_in,
        const uint16_t * __restrict__ weights_bf16,
        uint32_t idx,
        float inv_rms) {
    const float normalized = makepad_cuda_bf16_round(row_in[idx] * inv_rms);
    const float weight = makepad_cuda_bf16_to_f32(weights_bf16[idx]);
    return makepad_cuda_bf16_round(normalized * weight);
}

static __global__ void makepad_cuda_qkv_norm_rope_cache_f32_kernel(
        const float * __restrict__ qkv,
        const uint16_t * __restrict__ q_weights_bf16,
        const uint16_t * __restrict__ k_weights_bf16,
        float * __restrict__ q_out,
        uint16_t * __restrict__ key_cache,
        uint16_t * __restrict__ value_cache,
        uint32_t q_head_count,
        uint32_t k_head_count,
        uint32_t head_dim,
        uint32_t q_offset,
        uint32_t k_offset,
        uint32_t v_offset,
        uint32_t rotary_dim,
        float base,
        uint32_t position,
        float eps,
        uint32_t max_tokens,
        uint32_t slot) {
    const uint32_t row = blockIdx.x;
    const uint32_t total_rows = q_head_count + 2u * k_head_count;
    if (row >= total_rows) {
        return;
    }

    const bool is_q = row < q_head_count;
    const bool is_k = row >= q_head_count && row < q_head_count + k_head_count;
    const uint32_t local_row = is_q ? row : (row - q_head_count) % k_head_count;
    const uint32_t source_offset = is_q ? q_offset : (is_k ? k_offset : v_offset);
    const float * row_in = qkv + source_offset + local_row * head_dim;

    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < head_dim; idx += blockDim.x) {
        const float value = row_in[idx];
        sum += value * value;
    }
    sum = makepad_cuda_block_reduce_sum(sum);

    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(head_dim) + eps);
    }
    __syncthreads();

    const uint32_t half = head_dim / 2u;
    const uint32_t rotary_pairs = rotary_dim / 2u;
    const uint32_t cache_head_base = local_row * max_tokens * head_dim;
    const uint32_t cache_key_base = cache_head_base + slot * head_dim;

    for (uint32_t idx = threadIdx.x; idx < head_dim; idx += blockDim.x) {
        if (!is_q && !is_k) {
            value_cache[cache_head_base + idx * max_tokens + slot] =
                makepad_cuda_f32_to_bf16_bits(row_in[idx] * inv_rms);
            continue;
        }

        const uint16_t * weights_bf16 = is_q ? q_weights_bf16 : k_weights_bf16;
        float out_value = 0.0f;
        if (idx < rotary_pairs) {
            const float exponent = (2.0f * static_cast<float>(idx)) / static_cast<float>(head_dim);
            const float inv_freq = powf(base, -exponent);
            const float theta = static_cast<float>(position) * inv_freq;
            const float cos_theta = cosf(theta);
            const float sin_theta = sinf(theta);
            const float left = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
            const float right = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, half + idx, inv_rms);
            out_value = makepad_cuda_bf16_round(left * cos_theta - right * sin_theta);
        } else if (idx >= half && idx < half + rotary_pairs) {
            const uint32_t pair = idx - half;
            const float exponent = (2.0f * static_cast<float>(pair)) / static_cast<float>(head_dim);
            const float inv_freq = powf(base, -exponent);
            const float theta = static_cast<float>(position) * inv_freq;
            const float cos_theta = cosf(theta);
            const float sin_theta = sinf(theta);
            const float left = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, pair, inv_rms);
            const float right = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
            out_value = makepad_cuda_bf16_round(left * sin_theta + right * cos_theta);
        } else {
            out_value = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
        }

        if (is_q) {
            q_out[local_row * head_dim + idx] = out_value;
        } else {
            key_cache[cache_key_base + idx] = makepad_cuda_f32_to_bf16_bits(out_value);
        }
    }
}

static __global__ void makepad_cuda_qkv_norm_rope_cache_f32_device_u32_kernel(
        const float * __restrict__ qkv,
        const uint16_t * __restrict__ q_weights_bf16,
        const uint16_t * __restrict__ k_weights_bf16,
        float * __restrict__ q_out,
        uint16_t * __restrict__ key_cache,
        uint16_t * __restrict__ value_cache,
        uint32_t q_head_count,
        uint32_t k_head_count,
        uint32_t head_dim,
        uint32_t q_offset,
        uint32_t k_offset,
        uint32_t v_offset,
        uint32_t rotary_dim,
        float base,
        const uint32_t * __restrict__ position_device_u32,
        float eps,
        uint32_t max_tokens) {
    const uint32_t position = *position_device_u32;
    if (max_tokens == 0) {
        return;
    }
    const uint32_t slot = position % max_tokens;

    const uint32_t row = blockIdx.x;
    const uint32_t total_rows = q_head_count + 2u * k_head_count;
    if (row >= total_rows) {
        return;
    }

    const bool is_q = row < q_head_count;
    const bool is_k = row >= q_head_count && row < q_head_count + k_head_count;
    const uint32_t local_row = is_q ? row : (row - q_head_count) % k_head_count;
    const uint32_t source_offset = is_q ? q_offset : (is_k ? k_offset : v_offset);
    const float * row_in = qkv + source_offset + local_row * head_dim;

    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < head_dim; idx += blockDim.x) {
        const float value = row_in[idx];
        sum += value * value;
    }
    sum = makepad_cuda_block_reduce_sum(sum);

    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(head_dim) + eps);
    }
    __syncthreads();

    const uint32_t half = head_dim / 2u;
    const uint32_t rotary_pairs = rotary_dim / 2u;
    const uint32_t cache_head_base = local_row * max_tokens * head_dim;
    const uint32_t cache_key_base = cache_head_base + slot * head_dim;

    for (uint32_t idx = threadIdx.x; idx < head_dim; idx += blockDim.x) {
        if (!is_q && !is_k) {
            value_cache[cache_head_base + idx * max_tokens + slot] =
                makepad_cuda_f32_to_bf16_bits(row_in[idx] * inv_rms);
            continue;
        }

        const uint16_t * weights_bf16 = is_q ? q_weights_bf16 : k_weights_bf16;
        float out_value = 0.0f;
        if (idx < rotary_pairs) {
            const float exponent = (2.0f * static_cast<float>(idx)) / static_cast<float>(head_dim);
            const float inv_freq = powf(base, -exponent);
            const float theta = static_cast<float>(position) * inv_freq;
            const float cos_theta = cosf(theta);
            const float sin_theta = sinf(theta);
            const float left = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
            const float right = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, half + idx, inv_rms);
            out_value = makepad_cuda_bf16_round(left * cos_theta - right * sin_theta);
        } else if (idx >= half && idx < half + rotary_pairs) {
            const uint32_t pair = idx - half;
            const float exponent = (2.0f * static_cast<float>(pair)) / static_cast<float>(head_dim);
            const float inv_freq = powf(base, -exponent);
            const float theta = static_cast<float>(position) * inv_freq;
            const float cos_theta = cosf(theta);
            const float sin_theta = sinf(theta);
            const float left = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, pair, inv_rms);
            const float right = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
            out_value = makepad_cuda_bf16_round(left * sin_theta + right * cos_theta);
        } else {
            out_value = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
        }

        if (is_q) {
            q_out[local_row * head_dim + idx] = out_value;
        } else {
            key_cache[cache_key_base + idx] = makepad_cuda_f32_to_bf16_bits(out_value);
        }
    }
}

static __global__ void makepad_cuda_qkv_norm_rope_cache_rows_f32_kernel(
        const float * __restrict__ qkv,
        const uint16_t * __restrict__ q_weights_bf16,
        const uint16_t * __restrict__ k_weights_bf16,
        float * __restrict__ q_out,
        uint16_t * __restrict__ key_cache,
        uint16_t * __restrict__ value_cache,
        uint32_t q_head_count,
        uint32_t k_head_count,
        uint32_t head_dim,
        uint32_t qkv_row_stride,
        uint32_t q_out_row_stride,
        uint32_t q_offset,
        uint32_t k_offset,
        uint32_t v_offset,
        uint32_t rotary_dim,
        float base,
        uint32_t start_position,
        float eps,
        uint32_t max_tokens,
        uint32_t start_slot,
        uint32_t row_count) {
    const uint32_t rows_per_token = q_head_count + 2u * k_head_count;
    const uint32_t row_index = blockIdx.x;
    if (row_index >= row_count * rows_per_token) {
        return;
    }

    const uint32_t token_idx = row_index / rows_per_token;
    const uint32_t row = row_index % rows_per_token;
    const uint32_t position = start_position + token_idx;
    const uint32_t slot = (start_slot + token_idx) % max_tokens;

    const bool is_q = row < q_head_count;
    const bool is_k = row >= q_head_count && row < q_head_count + k_head_count;
    const uint32_t local_row = is_q ? row : (row - q_head_count) % k_head_count;
    const uint32_t source_offset = is_q ? q_offset : (is_k ? k_offset : v_offset);
    const float * row_in = qkv + token_idx * qkv_row_stride + source_offset + local_row * head_dim;

    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < head_dim; idx += blockDim.x) {
        const float value = row_in[idx];
        sum += value * value;
    }
    sum = makepad_cuda_block_reduce_sum(sum);

    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(head_dim) + eps);
    }
    __syncthreads();

    const uint32_t half = head_dim / 2u;
    const uint32_t rotary_pairs = rotary_dim / 2u;
    const uint32_t cache_head_base = local_row * max_tokens * head_dim;
    const uint32_t cache_key_base = cache_head_base + slot * head_dim;

    for (uint32_t idx = threadIdx.x; idx < head_dim; idx += blockDim.x) {
        if (!is_q && !is_k) {
            value_cache[cache_head_base + idx * max_tokens + slot] =
                makepad_cuda_f32_to_bf16_bits(row_in[idx] * inv_rms);
            continue;
        }

        const uint16_t * weights_bf16 = is_q ? q_weights_bf16 : k_weights_bf16;
        float out_value = 0.0f;
        if (idx < rotary_pairs) {
            const float exponent = (2.0f * static_cast<float>(idx)) / static_cast<float>(head_dim);
            const float inv_freq = powf(base, -exponent);
            const float theta = static_cast<float>(position) * inv_freq;
            const float cos_theta = cosf(theta);
            const float sin_theta = sinf(theta);
            const float left = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
            const float right = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, half + idx, inv_rms);
            out_value = makepad_cuda_bf16_round(left * cos_theta - right * sin_theta);
        } else if (idx >= half && idx < half + rotary_pairs) {
            const uint32_t pair = idx - half;
            const float exponent = (2.0f * static_cast<float>(pair)) / static_cast<float>(head_dim);
            const float inv_freq = powf(base, -exponent);
            const float theta = static_cast<float>(position) * inv_freq;
            const float cos_theta = cosf(theta);
            const float sin_theta = sinf(theta);
            const float left = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, pair, inv_rms);
            const float right = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
            out_value = makepad_cuda_bf16_round(left * sin_theta + right * cos_theta);
        } else {
            out_value = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
        }

        if (is_q) {
            q_out[token_idx * q_out_row_stride + local_row * head_dim + idx] = out_value;
        } else {
            key_cache[cache_key_base + idx] = makepad_cuda_f32_to_bf16_bits(out_value);
        }
    }
}

static __global__ void makepad_cuda_qkv_norm_rope_cache_rows_f32_device_u32_kernel(
        const float * __restrict__ qkv,
        const uint16_t * __restrict__ q_weights_bf16,
        const uint16_t * __restrict__ k_weights_bf16,
        float * __restrict__ q_out,
        uint16_t * __restrict__ key_cache,
        uint16_t * __restrict__ value_cache,
        uint32_t q_head_count,
        uint32_t k_head_count,
        uint32_t head_dim,
        uint32_t qkv_row_stride,
        uint32_t q_out_row_stride,
        uint32_t q_offset,
        uint32_t k_offset,
        uint32_t v_offset,
        uint32_t rotary_dim,
        float base,
        const uint32_t * __restrict__ start_position_device_u32,
        float eps,
        uint32_t max_tokens,
        const uint32_t * __restrict__ start_slot_device_u32,
        uint32_t row_count) {
    const uint32_t start_position = *start_position_device_u32;
    const uint32_t start_slot = *start_slot_device_u32;
    const uint32_t rows_per_token = q_head_count + 2u * k_head_count;
    const uint32_t row_index = blockIdx.x;
    if (row_index >= row_count * rows_per_token) {
        return;
    }

    const uint32_t token_idx = row_index / rows_per_token;
    const uint32_t row = row_index % rows_per_token;
    const uint32_t position = start_position + token_idx;
    const uint32_t slot = (start_slot + token_idx) % max_tokens;

    const bool is_q = row < q_head_count;
    const bool is_k = row >= q_head_count && row < q_head_count + k_head_count;
    const uint32_t local_row = is_q ? row : (row - q_head_count) % k_head_count;
    const uint32_t source_offset = is_q ? q_offset : (is_k ? k_offset : v_offset);
    const float * row_in = qkv + token_idx * qkv_row_stride + source_offset + local_row * head_dim;

    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < head_dim; idx += blockDim.x) {
        const float value = row_in[idx];
        sum += value * value;
    }
    sum = makepad_cuda_block_reduce_sum(sum);

    __shared__ float inv_rms;
    if (threadIdx.x == 0) {
        inv_rms = rsqrtf(sum / static_cast<float>(head_dim) + eps);
    }
    __syncthreads();

    const uint32_t half = head_dim / 2u;
    const uint32_t rotary_pairs = rotary_dim / 2u;
    const uint32_t cache_head_base = local_row * max_tokens * head_dim;
    const uint32_t cache_key_base = cache_head_base + slot * head_dim;

    for (uint32_t idx = threadIdx.x; idx < head_dim; idx += blockDim.x) {
        if (!is_q && !is_k) {
            value_cache[cache_head_base + idx * max_tokens + slot] =
                makepad_cuda_f32_to_bf16_bits(row_in[idx] * inv_rms);
            continue;
        }

        const uint16_t * weights_bf16 = is_q ? q_weights_bf16 : k_weights_bf16;
        float out_value = 0.0f;
        if (idx < rotary_pairs) {
            const float exponent = (2.0f * static_cast<float>(idx)) / static_cast<float>(head_dim);
            const float inv_freq = powf(base, -exponent);
            const float theta = static_cast<float>(position) * inv_freq;
            const float cos_theta = cosf(theta);
            const float sin_theta = sinf(theta);
            const float left = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
            const float right = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, half + idx, inv_rms);
            out_value = makepad_cuda_bf16_round(left * cos_theta - right * sin_theta);
        } else if (idx >= half && idx < half + rotary_pairs) {
            const uint32_t pair = idx - half;
            const float exponent = (2.0f * static_cast<float>(pair)) / static_cast<float>(head_dim);
            const float inv_freq = powf(base, -exponent);
            const float theta = static_cast<float>(position) * inv_freq;
            const float cos_theta = cosf(theta);
            const float sin_theta = sinf(theta);
            const float left = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, pair, inv_rms);
            const float right = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
            out_value = makepad_cuda_bf16_round(left * sin_theta + right * cos_theta);
        } else {
            out_value = makepad_cuda_weighted_norm_f32(row_in, weights_bf16, idx, inv_rms);
        }

        if (is_q) {
            q_out[token_idx * q_out_row_stride + local_row * head_dim + idx] = out_value;
        } else {
            key_cache[cache_key_base + idx] = makepad_cuda_f32_to_bf16_bits(out_value);
        }
    }
}

static __global__ void makepad_cuda_attention_logits_seq_f32_kernel(
        const float * __restrict__ q,
        const uint16_t * __restrict__ key_cache,
        float * __restrict__ logits,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t seq_len,
        uint32_t start_slot,
        uint32_t capacity,
        uint32_t logits_row_stride) {
    const uint32_t q_head = blockIdx.x;
    const uint32_t token = blockIdx.y;
    if (q_head >= q_head_count || token >= seq_len) {
        return;
    }
    const uint32_t kv_head = q_head / q_heads_per_kv;
    const uint32_t slot = (start_slot + token) % capacity;
    const float * q_row = q + q_head * head_dim;
    const uint16_t * k_row = key_cache + kv_head * kv_row_stride + slot * head_dim;
    float sum = 0.0f;
    for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
        sum += q_row[dim] * makepad_cuda_bf16_to_f32(k_row[dim]);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        logits[q_head * logits_row_stride + token] = makepad_cuda_bf16_round(sum);
    }
}

template <uint32_t tokens_per_block>
static __global__ void makepad_cuda_attention_logits_seq_f32_device_u32_kernel(
        const float * __restrict__ q,
        const uint16_t * __restrict__ key_cache,
        float * __restrict__ logits,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * __restrict__ seq_len_device_u32,
        const uint32_t * __restrict__ start_slot_device_u32,
        uint32_t capacity,
        uint32_t logits_row_stride) {
    const uint32_t q_head = blockIdx.x;
    if (q_head >= q_head_count) {
        return;
    }
    const uint32_t seq_len = *seq_len_device_u32;
    const uint32_t start_slot = *start_slot_device_u32;
    const uint32_t token_base = blockIdx.y * tokens_per_block;
    if (token_base >= seq_len) {
        return;
    }
    const uint32_t kv_head = q_head / q_heads_per_kv;
    const float * q_row = q + q_head * head_dim;
    const uint16_t * key_row = key_cache + kv_head * kv_row_stride;
    const bool cache_q = head_dim <= blockDim.x;
    __shared__ float shared_q[256];
    if (cache_q && threadIdx.x < head_dim) {
        shared_q[threadIdx.x] = q_row[threadIdx.x];
    }
    __syncthreads();

    for (uint32_t token_offset = 0; token_offset < tokens_per_block; ++token_offset) {
        const uint32_t token = token_base + token_offset;
        if (token >= seq_len) {
            break;
        }
        const uint32_t slot = (start_slot + token) % capacity;
        const uint16_t * k_row = key_row + slot * head_dim;
        float sum = 0.0f;
        if (cache_q) {
            if (threadIdx.x < head_dim) {
                sum = shared_q[threadIdx.x] * makepad_cuda_bf16_to_f32(k_row[threadIdx.x]);
            }
        } else {
            for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
                sum += q_row[dim] * makepad_cuda_bf16_to_f32(k_row[dim]);
            }
        }
        sum = makepad_cuda_block_reduce_sum(sum);
        if (threadIdx.x == 0) {
            logits[q_head * logits_row_stride + token] = makepad_cuda_bf16_round(sum);
        }
        __syncthreads();
    }
}

template <uint32_t tokens_per_block>
static __global__ void makepad_cuda_attention_logits_seq_f32_device_u32_vec2_kernel(
        const float * __restrict__ q,
        const uint16_t * __restrict__ key_cache,
        float * __restrict__ logits,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * __restrict__ seq_len_device_u32,
        const uint32_t * __restrict__ start_slot_device_u32,
        uint32_t capacity,
        uint32_t logits_row_stride) {
    const uint32_t q_head = blockIdx.x;
    if (q_head >= q_head_count) {
        return;
    }
    const uint32_t seq_len = *seq_len_device_u32;
    const uint32_t start_slot = *start_slot_device_u32;
    const uint32_t token_base = blockIdx.y * tokens_per_block;
    if (token_base >= seq_len) {
        return;
    }

    const uint32_t kv_head = q_head / q_heads_per_kv;
    const float2 * q_row = reinterpret_cast<const float2 *>(q + q_head * head_dim);
    const uint16_t * key_row = key_cache + kv_head * kv_row_stride;
    const uint32_t head_dim_pairs = head_dim >> 1;

    for (uint32_t token_offset = 0; token_offset < tokens_per_block; ++token_offset) {
        const uint32_t token = token_base + token_offset;
        if (token >= seq_len) {
            break;
        }
        const uint32_t slot = (start_slot + token) % capacity;
        const uint32_t * k_row_u32 = reinterpret_cast<const uint32_t *>(key_row + slot * head_dim);
        float sum = 0.0f;
        for (uint32_t pair = threadIdx.x; pair < head_dim_pairs; pair += blockDim.x) {
            const float2 qv = q_row[pair];
            const uint32_t packed = k_row_u32[pair];
            const float k0 = makepad_cuda_bf16_to_f32(static_cast<uint16_t>(packed));
            const float k1 = makepad_cuda_bf16_to_f32(static_cast<uint16_t>(packed >> 16));
            sum += qv.x * k0 + qv.y * k1;
        }
        sum = makepad_cuda_block_reduce_sum(sum);
        if (threadIdx.x == 0) {
            logits[q_head * logits_row_stride + token] = makepad_cuda_bf16_round(sum);
        }
        __syncthreads();
    }
}

static __global__ void makepad_cuda_softmax_rows_f32_kernel(
        const float * __restrict__ logits,
        float * __restrict__ probs,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t seq_len) {
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const float * row_in = logits + row * row_stride;
    float * row_out = probs + row * row_stride;

    float max_value = -CUDART_INF_F;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        const float value = row_in[idx];
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);
    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        sum += expf(row_in[idx] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        row_out[idx] = makepad_cuda_bf16_round(expf(row_in[idx] - shared_max) / shared_sum);
    }
}

static __global__ void makepad_cuda_softmax_rows_f32_device_u32_kernel(
        const float * __restrict__ logits,
        float * __restrict__ probs,
        uint32_t row_count,
        uint32_t row_stride,
        const uint32_t * __restrict__ seq_len_device_u32) {
    const uint32_t seq_len = *seq_len_device_u32;
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const float * row_in = logits + row * row_stride;
    float * row_out = probs + row * row_stride;

    float max_value = -CUDART_INF_F;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        const float value = row_in[idx];
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);
    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        sum += expf(row_in[idx] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        row_out[idx] = makepad_cuda_bf16_round(expf(row_in[idx] - shared_max) / shared_sum);
    }
}

static __global__ void makepad_cuda_softmax_rows_causal_f32_kernel(
        float * __restrict__ logits,
        uint32_t query_count,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t base_seq_len,
        uint32_t max_seq_len) {
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const uint32_t query_idx = row % query_count;
    const uint32_t seq_len = min(base_seq_len + query_idx + 1u, max_seq_len);
    float * row_values = logits + row * row_stride;

    float max_value = -CUDART_INF_F;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        const float value = makepad_cuda_bf16_round(row_values[idx]);
        row_values[idx] = value;
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);
    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        sum += expf(row_values[idx] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t idx = threadIdx.x; idx < max_seq_len; idx += blockDim.x) {
        row_values[idx] = idx < seq_len
            ? makepad_cuda_bf16_round(expf(row_values[idx] - shared_max) / shared_sum)
            : 0.0f;
    }
}

static __global__ void makepad_cuda_softmax_rows_causal_f32_device_u32_kernel(
        float * __restrict__ logits,
        uint32_t query_count,
        uint32_t row_count,
        uint32_t row_stride,
        const uint32_t * __restrict__ base_seq_len_device_u32,
        uint32_t max_seq_len) {
    const uint32_t base_seq_len = *base_seq_len_device_u32;
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const uint32_t query_idx = row % query_count;
    const uint32_t seq_len = min(base_seq_len + query_idx + 1u, max_seq_len);
    float * row_values = logits + row * row_stride;

    float max_value = -CUDART_INF_F;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        const float value = makepad_cuda_bf16_round(row_values[idx]);
        row_values[idx] = value;
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);
    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        sum += expf(row_values[idx] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t idx = threadIdx.x; idx < max_seq_len; idx += blockDim.x) {
        row_values[idx] = idx < seq_len
            ? makepad_cuda_bf16_round(expf(row_values[idx] - shared_max) / shared_sum)
            : 0.0f;
    }
}

static __global__ void makepad_cuda_softmax_rows_causal_bf16_kernel(
        const float * __restrict__ logits,
        uint16_t * __restrict__ probs,
        uint32_t query_count,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t base_seq_len,
        uint32_t max_seq_len) {
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const uint32_t query_idx = row % query_count;
    const uint32_t seq_len = min(base_seq_len + query_idx + 1u, max_seq_len);
    const float * row_logits = logits + row * row_stride;
    uint16_t * row_probs = probs + row * row_stride;

    float max_value = -CUDART_INF_F;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        const float value = makepad_cuda_bf16_round(row_logits[idx]);
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);
    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        sum += expf(makepad_cuda_bf16_round(row_logits[idx]) - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t idx = threadIdx.x; idx < max_seq_len; idx += blockDim.x) {
        const float value = idx < seq_len
            ? expf(makepad_cuda_bf16_round(row_logits[idx]) - shared_max) / shared_sum
            : 0.0f;
        row_probs[idx] = makepad_cuda_f32_to_bf16_bits(value);
    }
}

static __global__ void makepad_cuda_softmax_rows_causal_bf16_device_u32_kernel(
        const float * __restrict__ logits,
        uint16_t * __restrict__ probs,
        uint32_t query_count,
        uint32_t row_count,
        uint32_t row_stride,
        const uint32_t * __restrict__ base_seq_len_device_u32,
        uint32_t max_seq_len) {
    const uint32_t base_seq_len = *base_seq_len_device_u32;
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const uint32_t query_idx = row % query_count;
    const uint32_t seq_len = min(base_seq_len + query_idx + 1u, max_seq_len);
    const float * row_logits = logits + row * row_stride;
    uint16_t * row_probs = probs + row * row_stride;

    float max_value = -CUDART_INF_F;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        const float value = makepad_cuda_bf16_round(row_logits[idx]);
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);
    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < seq_len; idx += blockDim.x) {
        sum += expf(makepad_cuda_bf16_round(row_logits[idx]) - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t idx = threadIdx.x; idx < max_seq_len; idx += blockDim.x) {
        const float value = idx < seq_len
            ? expf(makepad_cuda_bf16_round(row_logits[idx]) - shared_max) / shared_sum
            : 0.0f;
        row_probs[idx] = makepad_cuda_f32_to_bf16_bits(value);
    }
}

static __global__ void makepad_cuda_softmax_rows_causal_vision_bf16_kernel(
        const float * __restrict__ logits,
        uint16_t * __restrict__ probs,
        uint32_t query_count,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t base_seq_len,
        uint32_t max_seq_len,
        uint32_t chunk_start_position,
        uint32_t vision_start_position,
        uint32_t vision_end_position) {
    const uint32_t row = blockIdx.x;
    if (row >= row_count) {
        return;
    }
    const uint32_t query_idx = row % query_count;
    const uint32_t query_position = chunk_start_position + query_idx;
    const bool query_is_vision =
        query_position >= vision_start_position && query_position <= vision_end_position;
    const uint32_t causal_seq_len = min(base_seq_len + query_idx + 1u, max_seq_len);
    const float * row_logits = logits + row * row_stride;
    uint16_t * row_probs = probs + row * row_stride;

    float max_value = -CUDART_INF_F;
    for (uint32_t idx = threadIdx.x; idx < max_seq_len; idx += blockDim.x) {
        const bool same_vision_group =
            query_is_vision && idx >= vision_start_position && idx <= vision_end_position;
        const bool allowed = idx < causal_seq_len || same_vision_group;
        if (allowed) {
            const float value = makepad_cuda_bf16_round(row_logits[idx]);
            max_value = value > max_value ? value : max_value;
        }
    }
    max_value = makepad_cuda_block_reduce_max(max_value);
    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t idx = threadIdx.x; idx < max_seq_len; idx += blockDim.x) {
        const bool same_vision_group =
            query_is_vision && idx >= vision_start_position && idx <= vision_end_position;
        const bool allowed = idx < causal_seq_len || same_vision_group;
        if (allowed) {
            sum += expf(makepad_cuda_bf16_round(row_logits[idx]) - shared_max);
        }
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t idx = threadIdx.x; idx < max_seq_len; idx += blockDim.x) {
        const bool same_vision_group =
            query_is_vision && idx >= vision_start_position && idx <= vision_end_position;
        const bool allowed = idx < causal_seq_len || same_vision_group;
        const float value = allowed
            ? expf(makepad_cuda_bf16_round(row_logits[idx]) - shared_max) / shared_sum
            : 0.0f;
        row_probs[idx] = makepad_cuda_f32_to_bf16_bits(value);
    }
}

static __global__ void makepad_cuda_attention_weighted_sum_f32_kernel(
        const float * __restrict__ probs,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t seq_len,
        uint32_t start_slot,
        uint32_t capacity,
        uint32_t probs_row_stride,
        uint32_t out_row_stride) {
    const uint32_t q_head = blockIdx.y;
    const uint32_t dim = blockIdx.x * blockDim.x + threadIdx.x;
    if (q_head >= q_head_count || dim >= head_dim) {
        return;
    }
    const uint32_t kv_head = q_head / q_heads_per_kv;
    const float * row_probs = probs + q_head * probs_row_stride;
    const uint16_t * value_row = value_cache + kv_head * kv_row_stride;
    float acc = 0.0f;
    for (uint32_t token = 0; token < seq_len; ++token) {
        const uint32_t slot = (start_slot + token) % capacity;
        const float value = makepad_cuda_bf16_to_f32(value_row[dim * capacity + slot]);
        acc += row_probs[token] * value;
    }
    out[q_head * out_row_stride + dim] = makepad_cuda_bf16_round(acc);
}

static __global__ void makepad_cuda_attention_weighted_sum_f32_device_u32_kernel(
        const float * __restrict__ probs,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * __restrict__ seq_len_device_u32,
        uint32_t capacity,
        uint32_t probs_row_stride,
        uint32_t out_row_stride) {
    const uint32_t seq_len = *seq_len_device_u32;
    const uint32_t q_head = blockIdx.y;
    const uint32_t dim = blockIdx.x * blockDim.x + threadIdx.x;
    if (q_head >= q_head_count || dim >= head_dim) {
        return;
    }
    const uint32_t kv_head = q_head / q_heads_per_kv;
    const float * row_probs = probs + q_head * probs_row_stride;
    const uint16_t * value_row = value_cache + kv_head * kv_row_stride;
    float acc = 0.0f;
    for (uint32_t token = 0; token < seq_len; ++token) {
        const float value = makepad_cuda_bf16_to_f32(value_row[dim * capacity + token]);
        acc += row_probs[token] * value;
    }
    out[q_head * out_row_stride + dim] = makepad_cuda_bf16_round(acc);
}

static __global__ void makepad_cuda_attention_softmax_weighted_sum_f32_kernel(
        const float * __restrict__ logits,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t seq_len,
        uint32_t start_slot,
        uint32_t capacity,
        uint32_t logits_row_stride,
        uint32_t out_row_stride) {
    extern __shared__ float shared_probs[];
    const uint32_t q_head = blockIdx.y;
    const uint32_t dim = blockIdx.x * blockDim.x + threadIdx.x;
    if (q_head >= q_head_count) {
        return;
    }
    const bool valid_dim = dim < head_dim;

    const float * row_logits = logits + q_head * logits_row_stride;
    float max_value = -CUDART_INF_F;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        const float value = row_logits[token];
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);

    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        sum += expf(row_logits[token] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        shared_probs[token] = makepad_cuda_bf16_round(expf(row_logits[token] - shared_max) / shared_sum);
    }
    __syncthreads();

    if (valid_dim) {
        const uint32_t kv_head = q_head / q_heads_per_kv;
        const uint16_t * value_row = value_cache + kv_head * kv_row_stride;
        float acc = 0.0f;
        for (uint32_t token = 0; token < seq_len; ++token) {
            const float prob = shared_probs[token];
            const uint32_t slot = (start_slot + token) % capacity;
            const float value = makepad_cuda_bf16_to_f32(value_row[dim * capacity + slot]);
            acc += prob * value;
        }
        out[q_head * out_row_stride + dim] = makepad_cuda_bf16_round(acc);
    }
}

static __global__ void makepad_cuda_attention_softmax_weighted_sum_f32_device_u32_kernel(
        const float * __restrict__ logits,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * __restrict__ seq_len_device_u32,
        const uint32_t * __restrict__ start_slot_device_u32,
        uint32_t capacity,
        uint32_t logits_row_stride,
        uint32_t out_row_stride) {
    extern __shared__ float shared_probs[];
    const uint32_t seq_len = *seq_len_device_u32;
    const uint32_t q_head = blockIdx.y;
    const uint32_t dim = blockIdx.x * blockDim.x + threadIdx.x;
    if (q_head >= q_head_count) {
        return;
    }
    const bool valid_dim = dim < head_dim;
    const uint32_t start_slot = *start_slot_device_u32;

    const float * row_logits = logits + q_head * logits_row_stride;
    float max_value = -CUDART_INF_F;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        const float value = row_logits[token];
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);

    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        sum += expf(row_logits[token] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        shared_probs[token] = makepad_cuda_bf16_round(expf(row_logits[token] - shared_max) / shared_sum);
    }
    __syncthreads();

    if (valid_dim) {
        const uint32_t kv_head = q_head / q_heads_per_kv;
        const uint16_t * value_row = value_cache + kv_head * kv_row_stride;
        float acc = 0.0f;
        for (uint32_t token = 0; token < seq_len; ++token) {
            const float prob = shared_probs[token];
            const uint32_t slot = (start_slot + token) % capacity;
            const float value = makepad_cuda_bf16_to_f32(value_row[dim * capacity + slot]);
            acc += prob * value;
        }
        out[q_head * out_row_stride + dim] = makepad_cuda_bf16_round(acc);
    }
}

static __global__ void makepad_cuda_attention_softmax_weighted_sum_pair_f32_device_u32_kernel(
        const float * __restrict__ logits,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t kv_head_count,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * __restrict__ seq_len_device_u32,
        const uint32_t * __restrict__ start_slot_device_u32,
        uint32_t capacity,
        uint32_t logits_row_stride,
        uint32_t out_row_stride) {
    extern __shared__ float shared_probs[];
    float * shared_probs_0 = shared_probs;
    float * shared_probs_1 = shared_probs + capacity;

    const uint32_t seq_len = *seq_len_device_u32;
    const uint32_t start_slot = *start_slot_device_u32;
    const uint32_t kv_head = blockIdx.y;
    const uint32_t dim = blockIdx.x * blockDim.x + threadIdx.x;
    if (kv_head >= kv_head_count) {
        return;
    }
    const bool valid_dim = dim < head_dim;
    const uint32_t q_head_0 = kv_head * 2;
    const uint32_t q_head_1 = q_head_0 + 1;
    const float * row_logits_0 = logits + q_head_0 * logits_row_stride;
    const float * row_logits_1 = logits + q_head_1 * logits_row_stride;

    float max_value_0 = -CUDART_INF_F;
    float max_value_1 = -CUDART_INF_F;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        const float value_0 = row_logits_0[token];
        const float value_1 = row_logits_1[token];
        max_value_0 = value_0 > max_value_0 ? value_0 : max_value_0;
        max_value_1 = value_1 > max_value_1 ? value_1 : max_value_1;
    }
    max_value_0 = makepad_cuda_block_reduce_max(max_value_0);
    max_value_1 = makepad_cuda_block_reduce_max(max_value_1);

    __shared__ float shared_max_0;
    __shared__ float shared_max_1;
    __shared__ float shared_sum_0;
    __shared__ float shared_sum_1;
    if (threadIdx.x == 0) {
        shared_max_0 = max_value_0;
        shared_max_1 = max_value_1;
    }
    __syncthreads();

    float sum_0 = 0.0f;
    float sum_1 = 0.0f;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        sum_0 += expf(row_logits_0[token] - shared_max_0);
        sum_1 += expf(row_logits_1[token] - shared_max_1);
    }
    sum_0 = makepad_cuda_block_reduce_sum(sum_0);
    sum_1 = makepad_cuda_block_reduce_sum(sum_1);
    if (threadIdx.x == 0) {
        shared_sum_0 = sum_0;
        shared_sum_1 = sum_1;
    }
    __syncthreads();

    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        shared_probs_0[token] = makepad_cuda_bf16_round(expf(row_logits_0[token] - shared_max_0) / shared_sum_0);
        shared_probs_1[token] = makepad_cuda_bf16_round(expf(row_logits_1[token] - shared_max_1) / shared_sum_1);
    }
    __syncthreads();

    if (valid_dim) {
        const uint16_t * value_row = value_cache + kv_head * kv_row_stride;
        float acc_0 = 0.0f;
        float acc_1 = 0.0f;
        for (uint32_t token = 0; token < seq_len; ++token) {
            const uint32_t slot = (start_slot + token) % capacity;
            const float value = makepad_cuda_bf16_to_f32(value_row[dim * capacity + slot]);
            acc_0 += shared_probs_0[token] * value;
            acc_1 += shared_probs_1[token] * value;
        }
        out[q_head_0 * out_row_stride + dim] = makepad_cuda_bf16_round(acc_0);
        out[q_head_1 * out_row_stride + dim] = makepad_cuda_bf16_round(acc_1);
    }
}

static __global__ void makepad_cuda_attention_weighted_sum_transposed_pair_f32_device_u32_kernel(
        const float * __restrict__ probs,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t kv_head_count,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * __restrict__ seq_len_device_u32,
        const uint32_t * __restrict__ start_slot_device_u32,
        uint32_t capacity,
        uint32_t probs_row_stride,
        uint32_t out_row_stride) {
    const uint32_t dim = blockIdx.x;
    const uint32_t kv_head = blockIdx.y;
    if (dim >= head_dim || kv_head >= kv_head_count) {
        return;
    }

    const uint32_t seq_len = *seq_len_device_u32;
    const uint32_t start_slot = *start_slot_device_u32;
    const uint32_t q_head_0 = kv_head * 2;
    const uint32_t q_head_1 = q_head_0 + 1;
    const float * row_probs_0 = probs + q_head_0 * probs_row_stride;
    const float * row_probs_1 = probs + q_head_1 * probs_row_stride;
    const uint16_t * value_row = value_cache + kv_head * kv_row_stride + dim * capacity;

    float acc_0 = 0.0f;
    float acc_1 = 0.0f;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        const uint32_t slot = (start_slot + token) % capacity;
        const float value = makepad_cuda_bf16_to_f32(value_row[slot]);
        acc_0 += row_probs_0[token] * value;
        acc_1 += row_probs_1[token] * value;
    }

    acc_0 = makepad_cuda_block_reduce_sum(acc_0);
    __syncthreads();
    acc_1 = makepad_cuda_block_reduce_sum(acc_1);
    __syncthreads();
    if (threadIdx.x == 0) {
        out[q_head_0 * out_row_stride + dim] = makepad_cuda_bf16_round(acc_0);
        out[q_head_1 * out_row_stride + dim] = makepad_cuda_bf16_round(acc_1);
    }
}

static __global__ void makepad_cuda_attention_weighted_sum_transposed_group8_f32_device_u32_kernel(
        const float * __restrict__ probs,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t kv_head_count,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * __restrict__ seq_len_device_u32,
        const uint32_t * __restrict__ start_slot_device_u32,
        uint32_t capacity,
        uint32_t probs_row_stride,
        uint32_t out_row_stride) {
    const uint32_t dim = blockIdx.x;
    const uint32_t kv_head = blockIdx.y;
    if (dim >= head_dim || kv_head >= kv_head_count) {
        return;
    }

    const uint32_t seq_len = *seq_len_device_u32;
    const uint32_t start_slot = *start_slot_device_u32;
    const uint32_t q_head_base = kv_head * 8;
    const float * row_probs[8];
#pragma unroll
    for (int i = 0; i < 8; ++i) {
        row_probs[i] = probs + (q_head_base + static_cast<uint32_t>(i)) * probs_row_stride;
    }
    const uint16_t * value_row = value_cache + kv_head * kv_row_stride + dim * capacity;

    float acc[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        const uint32_t slot = (start_slot + token) % capacity;
        const float value = makepad_cuda_bf16_to_f32(value_row[slot]);
#pragma unroll
        for (int i = 0; i < 8; ++i) {
            acc[i] += row_probs[i][token] * value;
        }
    }

#pragma unroll
    for (int i = 0; i < 8; ++i) {
        acc[i] = makepad_cuda_block_reduce_sum(acc[i]);
        __syncthreads();
    }
    if (threadIdx.x == 0) {
#pragma unroll
        for (int i = 0; i < 8; ++i) {
            out[(q_head_base + static_cast<uint32_t>(i)) * out_row_stride + dim] =
                makepad_cuda_bf16_round(acc[i]);
        }
    }
}

static __global__ void makepad_cuda_attention_seq_softmax_weighted_sum_f32_kernel(
        const float * __restrict__ q,
        const uint16_t * __restrict__ key_cache,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t seq_len,
        uint32_t start_slot,
        uint32_t capacity,
        uint32_t out_row_stride) {
    const uint32_t q_head = blockIdx.x;
    if (q_head >= q_head_count) {
        return;
    }

    extern __shared__ float shared_logits[];
    const uint32_t kv_head = q_head / q_heads_per_kv;
    const float * q_row = q + q_head * head_dim;
    const uint16_t * key_row = key_cache + kv_head * kv_row_stride;
    const uint16_t * value_row = value_cache + kv_head * kv_row_stride;

    for (uint32_t token = 0; token < seq_len; ++token) {
        const uint32_t slot = (start_slot + token) % capacity;
        const uint16_t * k_row = key_row + slot * head_dim;
        float sum = 0.0f;
        for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
            sum += q_row[dim] * makepad_cuda_bf16_to_f32(k_row[dim]);
        }
        sum = makepad_cuda_block_reduce_sum(sum);
        if (threadIdx.x == 0) {
            shared_logits[token] = makepad_cuda_bf16_round(sum);
        }
        __syncthreads();
    }

    float max_value = -CUDART_INF_F;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        const float value = shared_logits[token];
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);

    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        sum += expf(shared_logits[token] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        shared_logits[token] = makepad_cuda_bf16_round(expf(shared_logits[token] - shared_max) / shared_sum);
    }
    __syncthreads();

    for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
        float acc = 0.0f;
        for (uint32_t token = 0; token < seq_len; ++token) {
            const float prob = shared_logits[token];
            const uint32_t slot = (start_slot + token) % capacity;
            const float value = makepad_cuda_bf16_to_f32(value_row[dim * capacity + slot]);
            acc = makepad_cuda_bf16_round(acc + makepad_cuda_bf16_round(prob * value));
        }
        out[q_head * out_row_stride + dim] = acc;
    }
}

static __global__ void makepad_cuda_attention_seq_softmax_weighted_sum_f32_device_u32_kernel(
        const float * __restrict__ q,
        const uint16_t * __restrict__ key_cache,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * __restrict__ seq_len_device_u32,
        uint32_t capacity,
        uint32_t out_row_stride) {
    const uint32_t seq_len = *seq_len_device_u32;
    const uint32_t q_head = blockIdx.x;
    if (q_head >= q_head_count) {
        return;
    }

    extern __shared__ float shared_logits[];
    const uint32_t kv_head = q_head / q_heads_per_kv;
    const float * q_row = q + q_head * head_dim;
    const uint16_t * key_row = key_cache + kv_head * kv_row_stride;
    const uint16_t * value_row = value_cache + kv_head * kv_row_stride;
    const bool cache_q = head_dim <= blockDim.x;
    const float q_value = cache_q && threadIdx.x < head_dim ? q_row[threadIdx.x] : 0.0f;

    for (uint32_t token = 0; token < seq_len; ++token) {
        const uint16_t * k_row = key_row + token * head_dim;
        float sum = 0.0f;
        if (cache_q) {
            if (threadIdx.x < head_dim) {
                sum = q_value * makepad_cuda_bf16_to_f32(k_row[threadIdx.x]);
            }
        } else {
            for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
                sum += q_row[dim] * makepad_cuda_bf16_to_f32(k_row[dim]);
            }
        }
        sum = makepad_cuda_block_reduce_sum(sum);
        if (threadIdx.x == 0) {
            shared_logits[token] = makepad_cuda_bf16_round(sum);
        }
        __syncthreads();
    }

    float max_value = -CUDART_INF_F;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        const float value = shared_logits[token];
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);

    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        sum += expf(shared_logits[token] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        shared_logits[token] = makepad_cuda_bf16_round(expf(shared_logits[token] - shared_max) / shared_sum);
    }
    __syncthreads();

    for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
        float acc = 0.0f;
        for (uint32_t token = 0; token < seq_len; ++token) {
            const float prob = shared_logits[token];
            const float value = makepad_cuda_bf16_to_f32(value_row[dim * capacity + token]);
            acc = makepad_cuda_bf16_round(acc + makepad_cuda_bf16_round(prob * value));
        }
        out[q_head * out_row_stride + dim] = acc;
    }
}

static __global__ void makepad_cuda_attention_seq_softmax_weighted_sum_f32_device_u32_parallel_tokens_kernel(
        const float * __restrict__ q,
        const uint16_t * __restrict__ key_cache,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * __restrict__ seq_len_device_u32,
        uint32_t capacity,
        uint32_t out_row_stride) {
    const uint32_t seq_len = *seq_len_device_u32;
    const uint32_t q_head = blockIdx.x;
    if (q_head >= q_head_count) {
        return;
    }

    extern __shared__ float shared[];
    float * shared_logits = shared;
    float * shared_q = shared + capacity;

    const uint32_t kv_head = q_head / q_heads_per_kv;
    const float * q_row = q + q_head * head_dim;
    const uint16_t * key_row = key_cache + kv_head * kv_row_stride;
    const uint16_t * value_row = value_cache + kv_head * kv_row_stride;

    const uint32_t tid = threadIdx.x;
    const uint32_t lane = tid & 31u;
    const uint32_t group = tid >> 6;
    const uint32_t lane_in_group = tid & 63u;
    const uint32_t warp_in_group = (tid >> 5) & 1u;
    constexpr uint32_t group_size = 64;
    constexpr uint32_t groups_per_block = 256 / group_size;
    __shared__ float shared_group_partials[groups_per_block][2];

    for (uint32_t dim = tid; dim < head_dim; dim += blockDim.x) {
        shared_q[dim] = q_row[dim];
    }
    __syncthreads();

    for (uint32_t token_base = 0; token_base < seq_len; token_base += groups_per_block) {
        const uint32_t token = token_base + group;
        if (token < seq_len) {
            const uint16_t * k_row = key_row + token * head_dim;
            float sum = 0.0f;
            for (uint32_t dim = lane_in_group; dim < head_dim; dim += group_size) {
                sum += shared_q[dim] * makepad_cuda_bf16_to_f32(k_row[dim]);
            }
            sum = makepad_cuda_warp_reduce_sum(sum);
            if (lane == 0) {
                shared_group_partials[group][warp_in_group] = sum;
            }
        }
        __syncthreads();
        if (token < seq_len && warp_in_group == 0 && lane == 0) {
            shared_logits[token] = makepad_cuda_bf16_round(
                shared_group_partials[group][0] + shared_group_partials[group][1]);
        }
        __syncthreads();
    }

    float max_value = -CUDART_INF_F;
    for (uint32_t token = tid; token < seq_len; token += blockDim.x) {
        const float value = shared_logits[token];
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);

    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (tid == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t token = tid; token < seq_len; token += blockDim.x) {
        sum += expf(shared_logits[token] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (tid == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t token = tid; token < seq_len; token += blockDim.x) {
        shared_logits[token] = makepad_cuda_bf16_round(expf(shared_logits[token] - shared_max) / shared_sum);
    }
    __syncthreads();

    for (uint32_t dim = tid; dim < head_dim; dim += blockDim.x) {
        float acc = 0.0f;
        for (uint32_t token = 0; token < seq_len; ++token) {
            const float prob = shared_logits[token];
            const float value = makepad_cuda_bf16_to_f32(value_row[dim * capacity + token]);
            acc = makepad_cuda_bf16_round(acc + makepad_cuda_bf16_round(prob * value));
        }
        out[q_head * out_row_stride + dim] = acc;
    }
}

static __global__ void makepad_cuda_attention_seq_softmax_weighted_sum_rows_f32_kernel(
        const float * __restrict__ q,
        const uint16_t * __restrict__ key_cache,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t query_count,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t q_row_stride,
        uint32_t out_row_stride,
        uint32_t base_seq_len,
        uint32_t capacity) {
    const uint32_t q_head = blockIdx.x;
    const uint32_t query_idx = blockIdx.y;
    if (q_head >= q_head_count || query_idx >= query_count) {
        return;
    }

    const uint32_t query_seq_len = base_seq_len + query_idx + 1u;
    const uint32_t seq_len = min(query_seq_len, capacity);
    const bool wraps = query_seq_len > capacity;
    const uint32_t start_slot = wraps ? (query_seq_len - seq_len) % capacity : 0u;
    extern __shared__ float shared_logits[];
    const uint32_t kv_head = q_head / q_heads_per_kv;
    const float * q_row = q + query_idx * q_row_stride + q_head * head_dim;
    const uint16_t * key_row = key_cache + kv_head * kv_row_stride;
    const uint16_t * value_row = value_cache + kv_head * kv_row_stride;
    const bool cache_q = head_dim <= blockDim.x;
    const float q_value = cache_q && threadIdx.x < head_dim ? q_row[threadIdx.x] : 0.0f;

    for (uint32_t token = 0; token < seq_len; ++token) {
        const uint32_t slot = wraps ? (start_slot + token) % capacity : token;
        const uint16_t * k_row = key_row + slot * head_dim;
        float sum = 0.0f;
        if (cache_q) {
            if (threadIdx.x < head_dim) {
                sum = q_value * makepad_cuda_bf16_to_f32(k_row[threadIdx.x]);
            }
        } else {
            for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
                sum += q_row[dim] * makepad_cuda_bf16_to_f32(k_row[dim]);
            }
        }
        sum = makepad_cuda_block_reduce_sum(sum);
        if (threadIdx.x == 0) {
            shared_logits[token] = makepad_cuda_bf16_round(sum);
        }
        __syncthreads();
    }

    float max_value = -CUDART_INF_F;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        const float value = shared_logits[token];
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);

    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        sum += expf(shared_logits[token] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        shared_logits[token] = makepad_cuda_bf16_round(expf(shared_logits[token] - shared_max) / shared_sum);
    }
    __syncthreads();

    float * out_row = out + query_idx * out_row_stride + q_head * head_dim;
    for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
        float acc = 0.0f;
        for (uint32_t token = 0; token < seq_len; ++token) {
            const float prob = shared_logits[token];
            const uint32_t slot = wraps ? (start_slot + token) % capacity : token;
            const float value = makepad_cuda_bf16_to_f32(value_row[dim * capacity + slot]);
            acc = makepad_cuda_bf16_round(acc + makepad_cuda_bf16_round(prob * value));
        }
        out_row[dim] = acc;
    }
}

static __global__ void makepad_cuda_attention_seq_softmax_weighted_sum_rows_f32_device_u32_kernel(
        const float * __restrict__ q,
        const uint16_t * __restrict__ key_cache,
        const uint16_t * __restrict__ value_cache,
        float * __restrict__ out,
        uint32_t query_count,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t q_row_stride,
        uint32_t out_row_stride,
        const uint32_t * __restrict__ base_seq_len_device_u32,
        uint32_t capacity) {
    const uint32_t q_head = blockIdx.x;
    const uint32_t query_idx = blockIdx.y;
    if (q_head >= q_head_count || query_idx >= query_count) {
        return;
    }

    const uint32_t base_seq_len = *base_seq_len_device_u32;
    const uint32_t query_seq_len = base_seq_len + query_idx + 1u;
    const uint32_t seq_len = min(query_seq_len, capacity);
    const bool wraps = query_seq_len > capacity;
    const uint32_t start_slot = wraps ? (query_seq_len - seq_len) % capacity : 0u;
    extern __shared__ float shared_logits[];
    const uint32_t kv_head = q_head / q_heads_per_kv;
    const float * q_row = q + query_idx * q_row_stride + q_head * head_dim;
    const uint16_t * key_row = key_cache + kv_head * kv_row_stride;
    const uint16_t * value_row = value_cache + kv_head * kv_row_stride;
    const bool cache_q = head_dim <= blockDim.x;
    const float q_value = cache_q && threadIdx.x < head_dim ? q_row[threadIdx.x] : 0.0f;

    for (uint32_t token = 0; token < seq_len; ++token) {
        const uint32_t slot = wraps ? (start_slot + token) % capacity : token;
        const uint16_t * k_row = key_row + slot * head_dim;
        float sum = 0.0f;
        if (cache_q) {
            if (threadIdx.x < head_dim) {
                sum = q_value * makepad_cuda_bf16_to_f32(k_row[threadIdx.x]);
            }
        } else {
            for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
                sum += q_row[dim] * makepad_cuda_bf16_to_f32(k_row[dim]);
            }
        }
        sum = makepad_cuda_block_reduce_sum(sum);
        if (threadIdx.x == 0) {
            shared_logits[token] = makepad_cuda_bf16_round(sum);
        }
        __syncthreads();
    }

    float max_value = -CUDART_INF_F;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        const float value = shared_logits[token];
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);

    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        sum += expf(shared_logits[token] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        shared_logits[token] = makepad_cuda_bf16_round(expf(shared_logits[token] - shared_max) / shared_sum);
    }
    __syncthreads();

    float * out_row = out + query_idx * out_row_stride + q_head * head_dim;
    for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
        float acc = 0.0f;
        for (uint32_t token = 0; token < seq_len; ++token) {
            const float prob = shared_logits[token];
            const uint32_t slot = wraps ? (start_slot + token) % capacity : token;
            const float value = makepad_cuda_bf16_to_f32(value_row[dim * capacity + slot]);
            acc = makepad_cuda_bf16_round(acc + makepad_cuda_bf16_round(prob * value));
        }
        out_row[dim] = acc;
    }
}

static __global__ void makepad_cuda_flash_attn_f32_packed_kernel(
        const float * __restrict__ q,
        const float * __restrict__ k,
        const float * __restrict__ v,
        float * __restrict__ out,
        uint32_t seq_len,
        uint32_t num_heads,
        uint32_t head_dim,
        float scale) {
    const uint32_t head = blockIdx.x;
    const uint32_t query = blockIdx.y;
    if (head >= num_heads || query >= seq_len) {
        return;
    }

    extern __shared__ float shared_logits[];
    const uint32_t row_stride = num_heads * head_dim;
    const float * q_row = q + query * row_stride + head * head_dim;
    for (uint32_t token = 0; token < seq_len; ++token) {
        const float * k_row = k + token * row_stride + head * head_dim;
        float sum = 0.0f;
        for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
            sum += q_row[dim] * k_row[dim];
        }
        sum = makepad_cuda_block_reduce_sum(sum);
        if (threadIdx.x == 0) {
            shared_logits[token] = sum * scale;
        }
        __syncthreads();
    }

    float max_value = -CUDART_INF_F;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        const float value = shared_logits[token];
        max_value = value > max_value ? value : max_value;
    }
    max_value = makepad_cuda_block_reduce_max(max_value);

    __shared__ float shared_max;
    __shared__ float shared_sum;
    if (threadIdx.x == 0) {
        shared_max = max_value;
    }
    __syncthreads();

    float sum = 0.0f;
    for (uint32_t token = threadIdx.x; token < seq_len; token += blockDim.x) {
        sum += expf(shared_logits[token] - shared_max);
    }
    sum = makepad_cuda_block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        shared_sum = sum;
    }
    __syncthreads();

    float * out_row = out + query * row_stride + head * head_dim;
    for (uint32_t dim = threadIdx.x; dim < head_dim; dim += blockDim.x) {
        float acc = 0.0f;
        for (uint32_t token = 0; token < seq_len; ++token) {
            const float prob = expf(shared_logits[token] - shared_max) / shared_sum;
            const float value = v[token * row_stride + head * head_dim + dim];
            acc += prob * value;
        }
        out_row[dim] = acc;
    }
}

static __global__ void makepad_cuda_argmax_f32_kernel(
        const float * __restrict__ logits,
        uint32_t * __restrict__ out_index,
        uint32_t n) {
    float best_value = -CUDART_INF_F;
    uint32_t best_index = 0;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        const float value = logits[idx];
        if (value > best_value || (value == best_value && idx < best_index)) {
            best_value = value;
            best_index = idx;
        }
    }

    __shared__ float shared_values[1024];
    __shared__ uint32_t shared_indices[1024];
    shared_values[threadIdx.x] = best_value;
    shared_indices[threadIdx.x] = best_index;
    __syncthreads();

    for (uint32_t stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            const float other_value = shared_values[threadIdx.x + stride];
            const uint32_t other_index = shared_indices[threadIdx.x + stride];
            const float self_value = shared_values[threadIdx.x];
            const uint32_t self_index = shared_indices[threadIdx.x];
            if (other_value > self_value || (other_value == self_value && other_index < self_index)) {
                shared_values[threadIdx.x] = other_value;
                shared_indices[threadIdx.x] = other_index;
            }
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        *out_index = shared_indices[0];
    }
}

static __device__ __forceinline__ bool makepad_cuda_token_is_disallowed(
        uint32_t token_id,
        const uint32_t * __restrict__ disallowed_token_ids,
        uint32_t disallowed_count) {
    for (uint32_t index = 0; index < disallowed_count; ++index) {
        if (disallowed_token_ids[index] == token_id) {
            return true;
        }
    }
    return false;
}

static __global__ void makepad_cuda_mask_indices_f32_kernel(
        float * __restrict__ logits,
        const uint32_t * __restrict__ disallowed_token_ids,
        uint32_t disallowed_count,
        uint32_t n) {
    const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= disallowed_count) {
        return;
    }
    const uint32_t token_id = disallowed_token_ids[index];
    if (token_id < n) {
        logits[token_id] = -CUDART_INF_F;
    }
}

static __global__ void makepad_cuda_mask_indices_f32_device_u32_kernel(
        float * __restrict__ logits,
        const uint32_t * __restrict__ disallowed_token_ids,
        const uint32_t * __restrict__ disallowed_count_device_u32,
        uint32_t n) {
    const uint32_t disallowed_count = *disallowed_count_device_u32;
    const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= disallowed_count) {
        return;
    }
    const uint32_t token_id = disallowed_token_ids[index];
    if (token_id < n) {
        logits[token_id] = -CUDART_INF_F;
    }
}

static __device__ __forceinline__ bool makepad_cuda_argmax_candidate_is_better(
        float candidate_value,
        uint32_t candidate_index,
        float current_value,
        uint32_t current_index) {
    if (candidate_index == UINT32_MAX) {
        return false;
    }
    if (current_index == UINT32_MAX) {
        return true;
    }
    return candidate_value > current_value
        || (candidate_value == current_value && candidate_index < current_index);
}

static __device__ __forceinline__ void makepad_cuda_warp_reduce_argmax(
        float & value,
        uint32_t & index) {
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        const float other_value = __shfl_down_sync(0xffffffffu, value, offset);
        const uint32_t other_index = __shfl_down_sync(0xffffffffu, index, offset);
        if (makepad_cuda_argmax_candidate_is_better(
                other_value,
                other_index,
                value,
                index)) {
            value = other_value;
            index = other_index;
        }
    }
}

static __global__ void makepad_cuda_masked_argmax_f32_kernel(
        const float * __restrict__ logits,
        const uint32_t * __restrict__ disallowed_token_ids,
        uint32_t disallowed_count,
        uint32_t * __restrict__ out_index,
        uint32_t n) {
    float best_value = -CUDART_INF_F;
    uint32_t best_index = UINT32_MAX;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        if (makepad_cuda_token_is_disallowed(idx, disallowed_token_ids, disallowed_count)) {
            continue;
        }
        const float value = logits[idx];
        if (makepad_cuda_argmax_candidate_is_better(value, idx, best_value, best_index)) {
            best_value = value;
            best_index = idx;
        }
    }
    makepad_cuda_warp_reduce_argmax(best_value, best_index);

    __shared__ float shared_values[32];
    __shared__ uint32_t shared_indices[32];
    const uint32_t lane = threadIdx.x & 31u;
    const uint32_t warp = threadIdx.x >> 5;
    const uint32_t warp_count = (blockDim.x + 31u) >> 5;
    if (lane == 0) {
        shared_values[warp] = best_value;
        shared_indices[warp] = best_index;
    }
    __syncthreads();

    if (warp == 0) {
        best_value = lane < warp_count ? shared_values[lane] : -CUDART_INF_F;
        best_index = lane < warp_count ? shared_indices[lane] : UINT32_MAX;
        makepad_cuda_warp_reduce_argmax(best_value, best_index);
    }

    if (threadIdx.x == 0) {
        *out_index = best_index;
    }
}

static __global__ void makepad_cuda_masked_argmax_f32_device_u32_kernel(
        const float * __restrict__ logits,
        const uint32_t * __restrict__ disallowed_token_ids,
        const uint32_t * __restrict__ disallowed_count_device_u32,
        uint32_t * __restrict__ out_index,
        uint32_t n) {
    const uint32_t disallowed_count = *disallowed_count_device_u32;
    float best_value = -CUDART_INF_F;
    uint32_t best_index = UINT32_MAX;
    for (uint32_t idx = threadIdx.x; idx < n; idx += blockDim.x) {
        if (makepad_cuda_token_is_disallowed(idx, disallowed_token_ids, disallowed_count)) {
            continue;
        }
        const float value = logits[idx];
        if (makepad_cuda_argmax_candidate_is_better(value, idx, best_value, best_index)) {
            best_value = value;
            best_index = idx;
        }
    }
    makepad_cuda_warp_reduce_argmax(best_value, best_index);

    __shared__ float shared_values[32];
    __shared__ uint32_t shared_indices[32];
    const uint32_t lane = threadIdx.x & 31u;
    const uint32_t warp = threadIdx.x >> 5;
    const uint32_t warp_count = (blockDim.x + 31u) >> 5;
    if (lane == 0) {
        shared_values[warp] = best_value;
        shared_indices[warp] = best_index;
    }
    __syncthreads();

    if (warp == 0) {
        best_value = lane < warp_count ? shared_values[lane] : -CUDART_INF_F;
        best_index = lane < warp_count ? shared_indices[lane] : UINT32_MAX;
        makepad_cuda_warp_reduce_argmax(best_value, best_index);
    }

    if (threadIdx.x == 0) {
        *out_index = best_index;
    }
}

extern "C" cudaError_t makepad_cuda_quantize_q8_1_f32(
        const float * input_f32,
        uint8_t * output_q8_1_bytes,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0 || (n % 32) != 0) {
        return cudaErrorInvalidValue;
    }
    const uint32_t block_count = n / 32;
    makepad_cuda_quantize_q8_1_f32_kernel<<<block_count, 32, 0, stream>>>(
        input_f32,
        reinterpret_cast<block_q8_1 *>(output_q8_1_bytes),
        block_count);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_scale_f32_inplace(
        float * values,
        float scale,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaSuccess;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid((n + block.x - 1) / block.x, 1, 1);
    makepad_cuda_scale_f32_kernel<<<grid, block, 0, stream>>>(values, scale, n);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_scale_f32_inplace_device_f32_index(
        float * values,
        const float * scales,
        uint32_t scale_index,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaSuccess;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid((n + block.x - 1) / block.x, 1, 1);
    makepad_cuda_scale_f32_device_f32_index_kernel<<<grid, block, 0, stream>>>(
        values,
        scales,
        scale_index,
        n);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_f32_to_bf16(
        const float * input,
        uint16_t * output,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaSuccess;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid((n + block.x - 1) / block.x, 1, 1);
    makepad_cuda_f32_to_bf16_kernel<<<grid, block, 0, stream>>>(input, output, n);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_bf16_to_f32(
        const uint16_t * input,
        float * output,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaSuccess;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid((n + block.x - 1) / block.x, 1, 1);
    makepad_cuda_bf16_to_f32_kernel<<<grid, block, 0, stream>>>(input, output, n);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_skintokens_michelangelo_fourier_f32(
        const float * condition,
        float * output,
        uint32_t rows,
        cudaStream_t stream) {
    if (rows == 0) {
        return cudaSuccess;
    }
    if (rows > UINT32_MAX / 54u) {
        return cudaErrorInvalidValue;
    }
    const uint32_t n = rows * 54u;
    const dim3 block(256, 1, 1);
    const dim3 grid((n + block.x - 1) / block.x, 1, 1);
    makepad_cuda_skintokens_michelangelo_fourier_f32_kernel<<<grid, block, 0, stream>>>(
        condition,
        output,
        rows);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_add_f32(
        const float * left,
        const float * right,
        float * out,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaSuccess;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid((n + block.x - 1) / block.x, 1, 1);
    makepad_cuda_add_f32_kernel<<<grid, block, 0, stream>>>(left, right, out, n);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_copy_f32(
        const float * input,
        float * output,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaSuccess;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid((n + block.x - 1) / block.x, 1, 1);
    makepad_cuda_copy_f32_kernel<<<grid, block, 0, stream>>>(input, output, n);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_weighted_sum_rows_f32(
        const float * batched_inputs,
        const float * weights,
        float * output,
        uint32_t row_count,
        uint32_t input_count,
        cudaStream_t stream) {
    if (row_count == 0 || input_count == 0) {
        return cudaSuccess;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid((row_count + block.x - 1) / block.x, 1, 1);
    makepad_cuda_weighted_sum_rows_f32_kernel<<<grid, block, 0, stream>>>(
        batched_inputs,
        weights,
        output,
        row_count,
        input_count);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_weighted_sum_rows_grouped_f32(
        const float * batched_inputs,
        const float * weights,
        float * output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t input_count,
        cudaStream_t stream) {
    if (row_count == 0 || row_stride == 0 || input_count == 0) {
        return cudaSuccess;
    }
    const uint32_t total = row_count * row_stride;
    const dim3 block(256, 1, 1);
    const dim3 grid((total + block.x - 1) / block.x, 1, 1);
    makepad_cuda_weighted_sum_rows_grouped_f32_kernel<<<grid, block, 0, stream>>>(
        batched_inputs,
        weights,
        output,
        row_count,
        row_stride,
        input_count);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_add_scaled_rows_f32(
        const float * input,
        const float * scales,
        float * output,
        uint32_t row_count,
        uint32_t row_stride,
        cudaStream_t stream) {
    if (row_count == 0 || row_stride == 0) {
        return cudaSuccess;
    }
    const uint32_t total = row_count * row_stride;
    const dim3 block(256, 1, 1);
    const dim3 grid((total + block.x - 1) / block.x, 1, 1);
    makepad_cuda_add_scaled_rows_f32_kernel<<<grid, block, 0, stream>>>(
        input,
        scales,
        output,
        row_count,
        row_stride);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_add_scaled_rows_f32_indexed(
        const float * input,
        const float * scales,
        float * output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t scale_row_stride,
        uint32_t scale_column,
        cudaStream_t stream) {
    if (row_count == 0 || row_stride == 0 || scale_row_stride == 0) {
        return cudaSuccess;
    }
    const uint32_t total = row_count * row_stride;
    const dim3 block(256, 1, 1);
    const dim3 grid((total + block.x - 1) / block.x, 1, 1);
    makepad_cuda_add_scaled_rows_f32_indexed_kernel<<<grid, block, 0, stream>>>(
        input,
        scales,
        output,
        row_count,
        row_stride,
        scale_row_stride,
        scale_column);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_mul_f32(
        const float * left,
        const float * right,
        float * out,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaErrorInvalidValue;
    }
    const uint32_t block = 256;
    const uint32_t grid = (n + block - 1) / block;
    makepad_cuda_mul_f32_kernel<<<grid, block, 0, stream>>>(left, right, out, n);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_gelu_f32(
        const float * input,
        float * out,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaErrorInvalidValue;
    }
    const uint32_t block = 256;
    const uint32_t grid = (n + block - 1) / block;
    makepad_cuda_gelu_f32_kernel<<<grid, block, 0, stream>>>(input, out, n);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_geglu_split_f32(
        const float * gate_up,
        float * out,
        uint32_t n,
        uint32_t split_offset,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaSuccess;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid((n + block.x - 1) / block.x, 1, 1);
    makepad_cuda_geglu_split_f32_kernel<<<grid, block, 0, stream>>>(gate_up, out, n, split_offset);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_geglu_split_f32_rows(
        const float * gate_up,
        float * out,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        uint32_t split_offset,
        cudaStream_t stream) {
    if (row_count == 0 || n == 0) {
        return cudaSuccess;
    }
    const uint32_t total = row_count * n;
    const dim3 block(256, 1, 1);
    const dim3 grid((total + block.x - 1) / block.x, 1, 1);
    makepad_cuda_geglu_split_f32_rows_kernel<<<grid, block, 0, stream>>>(
        gate_up,
        out,
        row_count,
        row_stride,
        n,
        split_offset);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_rms_norm_row_weighted_f32(
        const float * input,
        const uint16_t * weights_bf16,
        float * output,
        uint32_t n,
        float eps,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaErrorInvalidValue;
    }
    const uint32_t block = n < 1024 ? 256 : 1024;
    makepad_cuda_rms_norm_row_weighted_f32_kernel<<<1, block, 0, stream>>>(input, weights_bf16, output, n, eps);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_rms_norm_row_weighted_f32_f32weights(
        const float * input,
        const float * weights_f32,
        float * output,
        uint32_t n,
        float eps,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaErrorInvalidValue;
    }
    const uint32_t block = n < 1024 ? 256 : 1024;
    makepad_cuda_rms_norm_row_weighted_f32_f32weights_kernel<<<1, block, 0, stream>>>(
        input, weights_f32, output, n, eps);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_rms_norm_row_weighted_f32_f32weights_precise(
        const float * input,
        const float * weights_f32,
        float * output,
        uint32_t n,
        float eps,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaErrorInvalidValue;
    }
    const uint32_t block = n < 1024 ? 256 : 1024;
    makepad_cuda_rms_norm_row_weighted_f32_f32weights_precise_kernel<<<1, block, 0, stream>>>(
        input, weights_f32, output, n, eps);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_rms_norm_rows_weighted_f32(
        const float * input,
        const uint16_t * weights_bf16,
        float * output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps,
        cudaStream_t stream) {
    if (row_count == 0 || n == 0 || row_stride < n) {
        return cudaErrorInvalidValue;
    }
    const uint32_t block = n < 1024 ? 256 : 1024;
    makepad_cuda_rms_norm_rows_weighted_f32_kernel<<<row_count, block, 0, stream>>>(
        input, weights_bf16, output, row_count, row_stride, n, eps);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_rms_norm_rows_weighted_f32_f32weights(
        const float * input,
        const float * weights_f32,
        float * output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps,
        cudaStream_t stream) {
    if (row_count == 0 || n == 0 || row_stride < n) {
        return cudaErrorInvalidValue;
    }
    const uint32_t block = n < 1024 ? 256 : 1024;
    makepad_cuda_rms_norm_rows_weighted_f32_f32weights_kernel<<<row_count, block, 0, stream>>>(
        input, weights_f32, output, row_count, row_stride, n, eps);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_rms_norm_rows_weighted_f32_f32weights_precise(
        const float * input,
        const float * weights_f32,
        float * output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps,
        cudaStream_t stream) {
    if (row_count == 0 || n == 0 || row_stride < n) {
        return cudaErrorInvalidValue;
    }
    const uint32_t block = n < 1024 ? 256 : 1024;
    makepad_cuda_rms_norm_rows_weighted_f32_f32weights_precise_kernel<<<row_count, block, 0, stream>>>(
        input, weights_f32, output, row_count, row_stride, n, eps);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_rms_norm_qwen3(
        const float * input,
        const float * weights_f32,
        float * output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps,
        cudaStream_t stream) {
    if (row_count == 0 || n == 0 || row_stride < n) {
        return cudaErrorInvalidValue;
    }
    const uint32_t block = n < 1024 ? 256 : 1024;
    makepad_cuda_rms_norm_qwen3_kernel<<<row_count, block, 0, stream>>>(
        input, weights_f32, output, row_count, row_stride, n, eps);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_rms_norm_rows_no_scale_f32(
        const float * input,
        float * output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps,
        cudaStream_t stream) {
    if (row_count == 0 || n == 0 || row_stride < n) {
        return cudaErrorInvalidValue;
    }
    makepad_cuda_rms_norm_rows_no_scale_f32_kernel<<<row_count, 256, 0, stream>>>(
        input, output, row_count, row_stride, n, eps);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_rms_norm_rows_no_scale_f32_precise(
        const float * input,
        float * output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t n,
        float eps,
        cudaStream_t stream) {
    if (row_count == 0 || n == 0 || row_stride < n) {
        return cudaErrorInvalidValue;
    }
    makepad_cuda_rms_norm_rows_no_scale_f32_precise_kernel<<<row_count, 256, 0, stream>>>(
        input, output, row_count, row_stride, n, eps);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_rope_rows_f32(
        const float * input,
        float * output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t head_dim,
        uint32_t rotary_dim,
        float base,
        uint32_t position,
        cudaStream_t stream) {
    if (row_count == 0 || head_dim == 0 || row_stride < head_dim || rotary_dim > head_dim || (rotary_dim & 1u) != 0u) {
        return cudaErrorInvalidValue;
    }
    const uint32_t total = row_count * row_stride;
    const dim3 block(256, 1, 1);
    const dim3 grid((total + block.x - 1) / block.x, 1, 1);
    makepad_cuda_rope_rows_f32_kernel<<<grid, block, 0, stream>>>(
        input, output, row_count, row_stride, head_dim, rotary_dim, base, position);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_rope_rows_f32_device_u32(
        const float * input,
        float * output,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t head_dim,
        uint32_t rotary_dim,
        float base,
        const uint32_t * position_device_u32,
        cudaStream_t stream) {
    if (row_count == 0 || head_dim == 0 || row_stride < head_dim || rotary_dim > head_dim || (rotary_dim & 1u) != 0u) {
        return cudaErrorInvalidValue;
    }
    const uint32_t total = row_count * row_stride;
    const dim3 block(256, 1, 1);
    const dim3 grid((total + block.x - 1) / block.x, 1, 1);
    makepad_cuda_rope_rows_f32_device_u32_kernel<<<grid, block, 0, stream>>>(
        input, output, row_count, row_stride, head_dim, rotary_dim, base, position_device_u32);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_kv_append_f32(
        const float * keys,
        const float * values,
        uint16_t * key_cache,
        uint16_t * value_cache,
        uint32_t kv_head_count,
        uint32_t head_dim,
        uint32_t max_tokens,
        uint32_t slot,
        cudaStream_t stream) {
    if (kv_head_count == 0 || head_dim == 0 || max_tokens == 0 || slot >= max_tokens) {
        return cudaErrorInvalidValue;
    }
    const uint32_t total = kv_head_count * head_dim;
    const dim3 block(256, 1, 1);
    const dim3 grid((total + block.x - 1) / block.x, 1, 1);
    makepad_cuda_kv_append_f32_kernel<<<grid, block, 0, stream>>>(
        keys, values, key_cache, value_cache, kv_head_count, head_dim, max_tokens, slot);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_kv_append_f32_device_u32(
        const float * keys,
        const float * values,
        uint16_t * key_cache,
        uint16_t * value_cache,
        uint32_t kv_head_count,
        uint32_t head_dim,
        uint32_t max_tokens,
        const uint32_t * slot_device_u32,
        cudaStream_t stream) {
    if (kv_head_count == 0 || head_dim == 0 || max_tokens == 0) {
        return cudaErrorInvalidValue;
    }
    const uint32_t total = kv_head_count * head_dim;
    const dim3 block(256, 1, 1);
    const dim3 grid((total + block.x - 1) / block.x, 1, 1);
    makepad_cuda_kv_append_f32_device_u32_kernel<<<grid, block, 0, stream>>>(
        keys, values, key_cache, value_cache, kv_head_count, head_dim, max_tokens, slot_device_u32);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_qkv_norm_rope_cache_f32(
        const float * qkv,
        const uint16_t * q_weights_bf16,
        const uint16_t * k_weights_bf16,
        float * q_out,
        uint16_t * key_cache,
        uint16_t * value_cache,
        uint32_t q_head_count,
        uint32_t k_head_count,
        uint32_t head_dim,
        uint32_t q_offset,
        uint32_t k_offset,
        uint32_t v_offset,
        uint32_t rotary_dim,
        float base,
        uint32_t position,
        float eps,
        uint32_t max_tokens,
        uint32_t slot,
        cudaStream_t stream) {
    if (q_head_count == 0 || k_head_count == 0 || head_dim == 0 || rotary_dim > head_dim ||
            (rotary_dim & 1u) != 0u || max_tokens == 0 || slot >= max_tokens) {
        return cudaErrorInvalidValue;
    }
    const uint32_t total_rows = q_head_count + 2u * k_head_count;
    makepad_cuda_qkv_norm_rope_cache_f32_kernel<<<total_rows, 256, 0, stream>>>(
        qkv,
        q_weights_bf16,
        k_weights_bf16,
        q_out,
        key_cache,
        value_cache,
        q_head_count,
        k_head_count,
        head_dim,
        q_offset,
        k_offset,
        v_offset,
        rotary_dim,
        base,
        position,
        eps,
        max_tokens,
        slot);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_qkv_norm_rope_cache_rows_f32(
        const float * qkv,
        const uint16_t * q_weights_bf16,
        const uint16_t * k_weights_bf16,
        float * q_out,
        uint16_t * key_cache,
        uint16_t * value_cache,
        uint32_t q_head_count,
        uint32_t k_head_count,
        uint32_t head_dim,
        uint32_t qkv_row_stride,
        uint32_t q_out_row_stride,
        uint32_t q_offset,
        uint32_t k_offset,
        uint32_t v_offset,
        uint32_t rotary_dim,
        float base,
        uint32_t start_position,
        float eps,
        uint32_t max_tokens,
        uint32_t start_slot,
        uint32_t row_count,
        cudaStream_t stream) {
    if (row_count == 0 || max_tokens == 0) {
        return cudaErrorInvalidValue;
    }
    const uint32_t total_rows = row_count * (q_head_count + 2u * k_head_count);
    makepad_cuda_qkv_norm_rope_cache_rows_f32_kernel<<<total_rows, 256, 0, stream>>>(
        qkv,
        q_weights_bf16,
        k_weights_bf16,
        q_out,
        key_cache,
        value_cache,
        q_head_count,
        k_head_count,
        head_dim,
        qkv_row_stride,
        q_out_row_stride,
        q_offset,
        k_offset,
        v_offset,
        rotary_dim,
        base,
        start_position,
        eps,
        max_tokens,
        start_slot,
        row_count);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_qkv_norm_rope_cache_f32_device_u32(
        const float * qkv,
        const uint16_t * q_weights_bf16,
        const uint16_t * k_weights_bf16,
        float * q_out,
        uint16_t * key_cache,
        uint16_t * value_cache,
        uint32_t q_head_count,
        uint32_t k_head_count,
        uint32_t head_dim,
        uint32_t q_offset,
        uint32_t k_offset,
        uint32_t v_offset,
        uint32_t rotary_dim,
        float base,
        const uint32_t * position_device_u32,
        float eps,
        uint32_t max_tokens,
        cudaStream_t stream) {
    if (q_head_count == 0 || k_head_count == 0 || head_dim == 0 || rotary_dim > head_dim ||
            (rotary_dim & 1u) != 0u || max_tokens == 0) {
        return cudaErrorInvalidValue;
    }
    const uint32_t total_rows = q_head_count + 2u * k_head_count;
    makepad_cuda_qkv_norm_rope_cache_f32_device_u32_kernel<<<total_rows, 256, 0, stream>>>(
        qkv,
        q_weights_bf16,
        k_weights_bf16,
        q_out,
        key_cache,
        value_cache,
        q_head_count,
        k_head_count,
        head_dim,
        q_offset,
        k_offset,
        v_offset,
        rotary_dim,
        base,
        position_device_u32,
        eps,
        max_tokens);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_qkv_norm_rope_cache_rows_f32_device_u32(
        const float * qkv,
        const uint16_t * q_weights_bf16,
        const uint16_t * k_weights_bf16,
        float * q_out,
        uint16_t * key_cache,
        uint16_t * value_cache,
        uint32_t q_head_count,
        uint32_t k_head_count,
        uint32_t head_dim,
        uint32_t qkv_row_stride,
        uint32_t q_out_row_stride,
        uint32_t q_offset,
        uint32_t k_offset,
        uint32_t v_offset,
        uint32_t rotary_dim,
        float base,
        const uint32_t * start_position_device_u32,
        float eps,
        uint32_t max_tokens,
        const uint32_t * start_slot_device_u32,
        uint32_t row_count,
        cudaStream_t stream) {
    if (row_count == 0 || max_tokens == 0) {
        return cudaErrorInvalidValue;
    }
    const uint32_t total_rows = row_count * (q_head_count + 2u * k_head_count);
    makepad_cuda_qkv_norm_rope_cache_rows_f32_device_u32_kernel<<<total_rows, 256, 0, stream>>>(
        qkv,
        q_weights_bf16,
        k_weights_bf16,
        q_out,
        key_cache,
        value_cache,
        q_head_count,
        k_head_count,
        head_dim,
        qkv_row_stride,
        q_out_row_stride,
        q_offset,
        k_offset,
        v_offset,
        rotary_dim,
        base,
        start_position_device_u32,
        eps,
        max_tokens,
        start_slot_device_u32,
        row_count);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_attention_logits_seq_f32(
        const float * q,
        const uint16_t * key_cache,
        float * logits,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t seq_len,
        uint32_t start_slot,
        uint32_t capacity,
        uint32_t logits_row_stride,
        cudaStream_t stream) {
    if (q_head_count == 0 || q_heads_per_kv == 0 || head_dim == 0 || seq_len == 0 || capacity == 0 || start_slot >= capacity || logits_row_stride < seq_len) {
        return cudaErrorInvalidValue;
    }
    const dim3 grid(q_head_count, seq_len, 1);
    makepad_cuda_attention_logits_seq_f32_kernel<<<grid, 256, 0, stream>>>(
        q, key_cache, logits, q_head_count, q_heads_per_kv, head_dim, kv_row_stride, seq_len, start_slot, capacity, logits_row_stride);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_attention_logits_seq_f32_device_u32(
        const float * q,
        const uint16_t * key_cache,
        float * logits,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * seq_len_device_u32,
        const uint32_t * start_slot_device_u32,
        uint32_t capacity,
        uint32_t logits_row_stride,
        cudaStream_t stream) {
    if (q_head_count == 0 || q_heads_per_kv == 0 || head_dim == 0 || capacity == 0 || logits_row_stride < capacity) {
        return cudaErrorInvalidValue;
    }
    constexpr uint32_t tokens_per_block = 4;
    const dim3 grid(q_head_count, (capacity + tokens_per_block - 1) / tokens_per_block, 1);
    makepad_cuda_attention_logits_seq_f32_device_u32_vec2_kernel<tokens_per_block><<<grid, 128, 0, stream>>>(
        q, key_cache, logits, q_head_count, q_heads_per_kv, head_dim, kv_row_stride, seq_len_device_u32, start_slot_device_u32, capacity, logits_row_stride);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_softmax_rows_f32(
        const float * logits,
        float * probs,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t seq_len,
        cudaStream_t stream) {
    if (row_count == 0 || row_stride < seq_len || seq_len == 0) {
        return cudaErrorInvalidValue;
    }
    makepad_cuda_softmax_rows_f32_kernel<<<row_count, 256, 0, stream>>>(
        logits, probs, row_count, row_stride, seq_len);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_softmax_rows_f32_device_u32(
        const float * logits,
        float * probs,
        uint32_t row_count,
        uint32_t row_stride,
        const uint32_t * seq_len_device_u32,
        cudaStream_t stream) {
    if (row_count == 0 || row_stride == 0) {
        return cudaErrorInvalidValue;
    }
    makepad_cuda_softmax_rows_f32_device_u32_kernel<<<row_count, 256, 0, stream>>>(
        logits, probs, row_count, row_stride, seq_len_device_u32);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_softmax_rows_causal_f32(
        float * logits,
        uint32_t query_count,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t base_seq_len,
        uint32_t max_seq_len,
        cudaStream_t stream) {
    if (query_count == 0 || row_count == 0 || row_stride < max_seq_len || max_seq_len == 0) {
        return cudaErrorInvalidValue;
    }
    makepad_cuda_softmax_rows_causal_f32_kernel<<<row_count, 256, 0, stream>>>(
        logits, query_count, row_count, row_stride, base_seq_len, max_seq_len);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_softmax_rows_causal_f32_device_u32(
        float * logits,
        uint32_t query_count,
        uint32_t row_count,
        uint32_t row_stride,
        const uint32_t * base_seq_len_device_u32,
        uint32_t max_seq_len,
        cudaStream_t stream) {
    if (query_count == 0 || row_count == 0 || row_stride < max_seq_len || max_seq_len == 0) {
        return cudaErrorInvalidValue;
    }
    makepad_cuda_softmax_rows_causal_f32_device_u32_kernel<<<row_count, 256, 0, stream>>>(
        logits, query_count, row_count, row_stride, base_seq_len_device_u32, max_seq_len);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_softmax_rows_causal_bf16(
        const float * logits,
        uint16_t * probs,
        uint32_t query_count,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t base_seq_len,
        uint32_t max_seq_len,
        cudaStream_t stream) {
    if (query_count == 0 || row_count == 0 || row_stride < max_seq_len || max_seq_len == 0) {
        return cudaErrorInvalidValue;
    }
    makepad_cuda_softmax_rows_causal_bf16_kernel<<<row_count, 256, 0, stream>>>(
        logits, probs, query_count, row_count, row_stride, base_seq_len, max_seq_len);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_softmax_rows_causal_bf16_device_u32(
        const float * logits,
        uint16_t * probs,
        uint32_t query_count,
        uint32_t row_count,
        uint32_t row_stride,
        const uint32_t * base_seq_len_device_u32,
        uint32_t max_seq_len,
        cudaStream_t stream) {
    if (query_count == 0 || row_count == 0 || row_stride < max_seq_len || max_seq_len == 0) {
        return cudaErrorInvalidValue;
    }
    makepad_cuda_softmax_rows_causal_bf16_device_u32_kernel<<<row_count, 256, 0, stream>>>(
        logits, probs, query_count, row_count, row_stride, base_seq_len_device_u32, max_seq_len);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_softmax_rows_causal_vision_bf16(
        const float * logits,
        uint16_t * probs,
        uint32_t query_count,
        uint32_t row_count,
        uint32_t row_stride,
        uint32_t base_seq_len,
        uint32_t max_seq_len,
        uint32_t chunk_start_position,
        uint32_t vision_start_position,
        uint32_t vision_end_position,
        cudaStream_t stream) {
    if (query_count == 0
            || row_count == 0
            || row_stride < max_seq_len
            || max_seq_len == 0
            || vision_start_position > vision_end_position) {
        return cudaErrorInvalidValue;
    }
    makepad_cuda_softmax_rows_causal_vision_bf16_kernel<<<row_count, 256, 0, stream>>>(
        logits,
        probs,
        query_count,
        row_count,
        row_stride,
        base_seq_len,
        max_seq_len,
        chunk_start_position,
        vision_start_position,
        vision_end_position);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_attention_weighted_sum_f32(
        const float * probs,
        const uint16_t * value_cache,
        float * out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t seq_len,
        uint32_t start_slot,
        uint32_t capacity,
        uint32_t probs_row_stride,
        uint32_t out_row_stride,
        cudaStream_t stream) {
    if (q_head_count == 0 || q_heads_per_kv == 0 || head_dim == 0 || seq_len == 0 || capacity == 0 || start_slot >= capacity || probs_row_stride < seq_len || out_row_stride < head_dim) {
        return cudaErrorInvalidValue;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid((head_dim + block.x - 1) / block.x, q_head_count, 1);
    makepad_cuda_attention_weighted_sum_f32_kernel<<<grid, block, 0, stream>>>(
        probs, value_cache, out, q_head_count, q_heads_per_kv, head_dim, kv_row_stride, seq_len, start_slot, capacity, probs_row_stride, out_row_stride);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_attention_softmax_weighted_sum_f32(
        const float * logits,
        const uint16_t * value_cache,
        float * out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t seq_len,
        uint32_t start_slot,
        uint32_t capacity,
        uint32_t logits_row_stride,
        uint32_t out_row_stride,
        cudaStream_t stream) {
    if (q_head_count == 0 || q_heads_per_kv == 0 || head_dim == 0 || seq_len == 0 || capacity == 0 || start_slot >= capacity || logits_row_stride < seq_len || out_row_stride < head_dim) {
        return cudaErrorInvalidValue;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid((head_dim + block.x - 1) / block.x, q_head_count, 1);
    const size_t shared_bytes = static_cast<size_t>(seq_len) * sizeof(float);
    makepad_cuda_attention_softmax_weighted_sum_f32_kernel<<<grid, block, shared_bytes, stream>>>(
        logits, value_cache, out, q_head_count, q_heads_per_kv, head_dim, kv_row_stride, seq_len, start_slot, capacity, logits_row_stride, out_row_stride);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_attention_weighted_sum_f32_device_u32(
        const float * probs,
        const uint16_t * value_cache,
        float * out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * seq_len_device_u32,
        uint32_t capacity,
        uint32_t probs_row_stride,
        uint32_t out_row_stride,
        cudaStream_t stream) {
    if (q_head_count == 0 || q_heads_per_kv == 0 || head_dim == 0 || capacity == 0 || probs_row_stride < capacity || out_row_stride < head_dim) {
        return cudaErrorInvalidValue;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid((head_dim + block.x - 1) / block.x, q_head_count, 1);
    makepad_cuda_attention_weighted_sum_f32_device_u32_kernel<<<grid, block, 0, stream>>>(
        probs, value_cache, out, q_head_count, q_heads_per_kv, head_dim, kv_row_stride, seq_len_device_u32, capacity, probs_row_stride, out_row_stride);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_attention_softmax_weighted_sum_f32_device_u32(
        const float * logits,
        const uint16_t * value_cache,
        float * out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * seq_len_device_u32,
        const uint32_t * start_slot_device_u32,
        uint32_t capacity,
        uint32_t logits_row_stride,
        uint32_t out_row_stride,
        cudaStream_t stream) {
    if (q_head_count == 0 || q_heads_per_kv == 0 || head_dim == 0 || capacity == 0 || logits_row_stride < capacity || out_row_stride < head_dim) {
        return cudaErrorInvalidValue;
    }
    const dim3 block(256, 1, 1);
    if (q_heads_per_kv == 2 && (q_head_count % 2) == 0) {
        makepad_cuda_softmax_rows_f32_device_u32_kernel<<<q_head_count, 256, 0, stream>>>(
            logits, const_cast<float *>(logits), q_head_count, logits_row_stride, seq_len_device_u32);
        cudaError_t status = cudaGetLastError();
        if (status != cudaSuccess) {
            return status;
        }
        const dim3 pv_block(128, 1, 1);
        const dim3 pv_grid(head_dim, q_head_count / 2, 1);
        makepad_cuda_attention_weighted_sum_transposed_pair_f32_device_u32_kernel<<<pv_grid, pv_block, 0, stream>>>(
            logits, value_cache, out, q_head_count / 2, head_dim, kv_row_stride, seq_len_device_u32, start_slot_device_u32, capacity, logits_row_stride, out_row_stride);
    } else if (q_heads_per_kv == 8 && (q_head_count % 8) == 0) {
        makepad_cuda_softmax_rows_f32_device_u32_kernel<<<q_head_count, 256, 0, stream>>>(
            logits, const_cast<float *>(logits), q_head_count, logits_row_stride, seq_len_device_u32);
        cudaError_t status = cudaGetLastError();
        if (status != cudaSuccess) {
            return status;
        }
        const dim3 pv_block(128, 1, 1);
        const dim3 pv_grid(head_dim, q_head_count / 8, 1);
        makepad_cuda_attention_weighted_sum_transposed_group8_f32_device_u32_kernel<<<pv_grid, pv_block, 0, stream>>>(
            logits, value_cache, out, q_head_count / 8, head_dim, kv_row_stride, seq_len_device_u32, start_slot_device_u32, capacity, logits_row_stride, out_row_stride);
    } else {
        const dim3 grid((head_dim + block.x - 1) / block.x, q_head_count, 1);
        const size_t shared_bytes = static_cast<size_t>(capacity) * sizeof(float);
        makepad_cuda_attention_softmax_weighted_sum_f32_device_u32_kernel<<<grid, block, shared_bytes, stream>>>(
            logits, value_cache, out, q_head_count, q_heads_per_kv, head_dim, kv_row_stride, seq_len_device_u32, start_slot_device_u32, capacity, logits_row_stride, out_row_stride);
    }
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_attention_seq_softmax_weighted_sum_f32(
        const float * q,
        const uint16_t * key_cache,
        const uint16_t * value_cache,
        float * out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t seq_len,
        uint32_t start_slot,
        uint32_t capacity,
        uint32_t out_row_stride,
        cudaStream_t stream) {
    if (q_head_count == 0 || q_heads_per_kv == 0 || head_dim == 0 || seq_len == 0 || capacity == 0 || start_slot >= capacity || out_row_stride < head_dim) {
        return cudaErrorInvalidValue;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid(q_head_count, 1, 1);
    const size_t shared_bytes = static_cast<size_t>(seq_len) * sizeof(float);
    makepad_cuda_attention_seq_softmax_weighted_sum_f32_kernel<<<grid, block, shared_bytes, stream>>>(
        q, key_cache, value_cache, out, q_head_count, q_heads_per_kv, head_dim, kv_row_stride, seq_len, start_slot, capacity, out_row_stride);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_attention_seq_softmax_weighted_sum_rows_f32(
        const float * q,
        const uint16_t * key_cache,
        const uint16_t * value_cache,
        float * out,
        uint32_t query_count,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t q_row_stride,
        uint32_t out_row_stride,
        uint32_t base_seq_len,
        uint32_t capacity,
        cudaStream_t stream) {
    if (query_count == 0 || q_head_count == 0 || head_dim == 0 || capacity == 0) {
        return cudaErrorInvalidValue;
    }
    const dim3 grid(q_head_count, query_count, 1);
    const dim3 block(256, 1, 1);
    const size_t shared_bytes = static_cast<size_t>(min(base_seq_len + query_count, capacity)) * sizeof(float);
    makepad_cuda_attention_seq_softmax_weighted_sum_rows_f32_kernel<<<grid, block, shared_bytes, stream>>>(
        q,
        key_cache,
        value_cache,
        out,
        query_count,
        q_head_count,
        q_heads_per_kv,
        head_dim,
        kv_row_stride,
        q_row_stride,
        out_row_stride,
        base_seq_len,
        capacity);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_attention_seq_softmax_weighted_sum_f32_device_u32(
        const float * q,
        const uint16_t * key_cache,
        const uint16_t * value_cache,
        float * out,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        const uint32_t * seq_len_device_u32,
        uint32_t capacity,
        uint32_t out_row_stride,
        cudaStream_t stream) {
    if (q_head_count == 0 || q_heads_per_kv == 0 || head_dim == 0 || capacity == 0 || out_row_stride < head_dim) {
        return cudaErrorInvalidValue;
    }
    const dim3 block(256, 1, 1);
    const dim3 grid(q_head_count, 1, 1);
    const size_t shared_bytes = static_cast<size_t>(capacity) * sizeof(float);
    makepad_cuda_attention_seq_softmax_weighted_sum_f32_device_u32_kernel<<<grid, block, shared_bytes, stream>>>(
        q, key_cache, value_cache, out, q_head_count, q_heads_per_kv, head_dim, kv_row_stride, seq_len_device_u32, capacity, out_row_stride);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_attention_seq_softmax_weighted_sum_rows_f32_device_u32(
        const float * q,
        const uint16_t * key_cache,
        const uint16_t * value_cache,
        float * out,
        uint32_t query_count,
        uint32_t q_head_count,
        uint32_t q_heads_per_kv,
        uint32_t head_dim,
        uint32_t kv_row_stride,
        uint32_t q_row_stride,
        uint32_t out_row_stride,
        const uint32_t * base_seq_len_device_u32,
        uint32_t capacity,
        cudaStream_t stream) {
    if (query_count == 0 || q_head_count == 0 || head_dim == 0 || capacity == 0) {
        return cudaErrorInvalidValue;
    }
    const dim3 grid(q_head_count, query_count, 1);
    const dim3 block(256, 1, 1);
    const size_t shared_bytes = static_cast<size_t>(capacity) * sizeof(float);
    makepad_cuda_attention_seq_softmax_weighted_sum_rows_f32_device_u32_kernel<<<grid, block, shared_bytes, stream>>>(
        q,
        key_cache,
        value_cache,
        out,
        query_count,
        q_head_count,
        q_heads_per_kv,
        head_dim,
        kv_row_stride,
        q_row_stride,
        out_row_stride,
        base_seq_len_device_u32,
        capacity);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_flash_attn_f32_packed(
        const float * q,
        const float * k,
        const float * v,
        float * out,
        uint32_t seq_len,
        uint32_t num_heads,
        uint32_t head_dim,
        float scale,
        cudaStream_t stream) {
    if (seq_len == 0 || num_heads == 0 || head_dim == 0) {
        return cudaErrorInvalidValue;
    }
    const dim3 grid(num_heads, seq_len, 1);
    const dim3 block(256, 1, 1);
    const size_t shared_bytes = static_cast<size_t>(seq_len) * sizeof(float);
    makepad_cuda_flash_attn_f32_packed_kernel<<<grid, block, shared_bytes, stream>>>(
        q, k, v, out, seq_len, num_heads, head_dim, scale);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_argmax_f32(
        const float * logits,
        uint32_t * out_index,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaErrorInvalidValue;
    }
    makepad_cuda_argmax_f32_kernel<<<1, 1024, 0, stream>>>(logits, out_index, n);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_masked_argmax_f32(
        const float * logits,
        const uint32_t * disallowed_token_ids,
        uint32_t disallowed_count,
        uint32_t * out_index,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaErrorInvalidValue;
    }
    if (disallowed_count == 0) {
        makepad_cuda_argmax_f32_kernel<<<1, 256, 0, stream>>>(logits, out_index, n);
        return cudaGetLastError();
    }
    makepad_cuda_masked_argmax_f32_kernel<<<1, 512, 0, stream>>>(
        logits,
        disallowed_token_ids,
        disallowed_count,
        out_index,
        n);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_masked_argmax_f32_device_u32(
        const float * logits,
        const uint32_t * disallowed_token_ids,
        const uint32_t * disallowed_count_device_u32,
        uint32_t * out_index,
        uint32_t n,
        cudaStream_t stream) {
    if (n == 0) {
        return cudaErrorInvalidValue;
    }
    makepad_cuda_masked_argmax_f32_device_u32_kernel<<<1, 512, 0, stream>>>(
        logits,
        disallowed_token_ids,
        disallowed_count_device_u32,
        out_index,
        n);
    return cudaGetLastError();
}