// Modulate kernel
// Applies per-channel affine and optional gating.

override WORKGROUP_SIZE: u32 = 256u;

struct Uniforms {
    num_tokens: u32,
    hidden_size: u32,
    scale_offset: u32,
    shift_offset: u32,
    gate_offset: u32,
    has_gate: u32,
    add_one: u32,
    _pad0: u32,
}

@group(0) @binding(0) var<uniform> u: Uniforms;
@group(0) @binding(1) var<storage, read> input: array<f32>;
@group(0) @binding(2) var<storage, read> mod_params: array<f32>;
@group(0) @binding(3) var<storage, read_write> output: array<f32>;

@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    let total = u.num_tokens * u.hidden_size;
    if (idx >= total) {
        return;
    }

    let dim = idx % u.hidden_size;
    let raw_scale = mod_params[u.scale_offset + dim];
    let shift = mod_params[u.shift_offset + dim];
    let scale = select(raw_scale, 1.0 + raw_scale, u.add_one != 0u);
    var value = input[idx] * scale + shift;
    if (u.has_gate != 0u) {
        let gate = mod_params[u.gate_offset + dim];
        value = value * gate;
    }
    output[idx] = value;
}