struct ds4_metal_args_norm {
    int32_t  ne00;
    int32_t  ne00_t;
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
    float    eps;
    int32_t  nef1[3];
    int32_t  nef2[3];
    int32_t  nef3[3];
    uint64_t nbf1[3];
    uint64_t nbf2[3];
    uint64_t nbf3[3];
};

// RMSNorm over one activation row, optionally fusing the learned weight
// multiply. DS4 calls this before attention, before the FFN, and for plain
// diagnostics that need normalized but unweighted rows.
template <typename T, short F>
kernel void kernel_rms_norm_fuse_impl(
        constant ds4_metal_args_norm & args,
        device const char * src0,
        device const char * src1_0,
        device const char * src1_1,
        device       char * dst,
        threadgroup float * shmem_f32 [[threadgroup(0)]],
        uint3   tgpig[[threadgroup_position_in_grid]],
        ushort3 tpitg[[thread_position_in_threadgroup]],
        ushort  sgitg[[simdgroup_index_in_threadgroup]],
        ushort  tiisg[[thread_index_in_simdgroup]],
        ushort3   ntg[[threads_per_threadgroup]]) {
    if (sgitg == 0) {
        shmem_f32[tiisg] = 0.0f;
    }

    const int i01 = tgpig.x;
    const int i02 = tgpig.y;
    const int i03 = tgpig.z;

    device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);

    device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
    device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]);

    float sumf = 0.0f;

    // parallel sum
    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
        sumf += dot(x[i00], x[i00]);
    }
    sumf = simd_sum(sumf);

    threadgroup_barrier(mem_flags::mem_threadgroup);

    if (tiisg == 0) {
        shmem_f32[sgitg] = sumf;
    }

    threadgroup_barrier(mem_flags::mem_threadgroup);

    sumf = shmem_f32[tiisg];
    sumf = simd_sum(sumf);

    const float mean  = sumf/args.ne00;
    const float scale = 1.0f/sqrt(mean + args.eps);

    device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
        if (F == 1) {
            y[i00] = (x[i00]*scale);
        }
        if (F == 2) {
            y[i00] = (x[i00]*scale)*f0[i00];
        }
        if (F == 3) {
            y[i00] = (x[i00]*scale)*f0[i00] + f1[i00];
        }
    }
}

typedef decltype(kernel_rms_norm_fuse_impl<float4, 1>) kernel_rms_norm_fuse_t;

// Host-visible RMSNorm variants: plain norm and norm multiplied by weight.
template [[host_name("kernel_rms_norm_f32_4")]]     kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 1>;
template [[host_name("kernel_rms_norm_mul_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 2>;

kernel void kernel_add_rms_norm_mul_f32_4(
        constant ds4_metal_args_norm & args,
        device const char * src0,
        device const char * src1,
        device const char * weight,
        device       char * sum_dst,
        device       char * norm_dst,
        threadgroup float * shmem_f32 [[threadgroup(0)]],
        uint3   tgpig[[threadgroup_position_in_grid]],
        ushort3 tpitg[[thread_position_in_threadgroup]],
        ushort  sgitg[[simdgroup_index_in_threadgroup]],
        ushort  tiisg[[thread_index_in_simdgroup]],
        ushort3 ntg[[threads_per_threadgroup]]) {
    if (sgitg == 0) shmem_f32[tiisg] = 0.0f;

    const int i01 = tgpig.x;
    const int i02 = tgpig.y;
    const int i03 = tgpig.z;
    device const float4 *a = (device const float4 *)
        (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
    device const float4 *b = (device const float4 *)
        (src1 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
    device const float4 *w = (device const float4 *)
        (weight + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
    device float4 *sum = (device float4 *)
        (sum_dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
    device float4 *norm = (device float4 *)
        (norm_dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);

    float sumf = 0.0f;
    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
        const float4 v = a[i00] + b[i00];
        sum[i00] = v;
        sumf += dot(v, v);
    }
    sumf = simd_sum(sumf);
    threadgroup_barrier(mem_flags::mem_threadgroup);
    if (tiisg == 0) shmem_f32[sgitg] = sumf;
    threadgroup_barrier(mem_flags::mem_threadgroup);
    sumf = simd_sum(shmem_f32[tiisg]);

    const float mean = sumf / args.ne00;
    const float scale = 1.0f / sqrt(mean + args.eps);
    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
        norm[i00] = (sum[i00] * scale) * w[i00];
    }
}

// RMSNorm reduction used when the following F16 matmul applies the row scale
// while staging its RHS tile.
kernel void kernel_rms_norm_scale_f32_4(
        constant ds4_metal_args_norm & args,
        device const char * src0,
        device       float * dst_scale,
        threadgroup float * shmem_f32 [[threadgroup(0)]],
        uint3   tgpig[[threadgroup_position_in_grid]],
        ushort3 tpitg[[thread_position_in_threadgroup]],
        ushort  sgitg[[simdgroup_index_in_threadgroup]],
        ushort  tiisg[[thread_index_in_simdgroup]],
        ushort3 ntg[[threads_per_threadgroup]]) {
    if (sgitg == 0) shmem_f32[tiisg] = 0.0f;

    const int i01 = tgpig.x;
    const int i02 = tgpig.y;
    const int i03 = tgpig.z;
    device const float4 *x = (device const float4 *)
        (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);

    float sumf = 0.0f;
    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
        sumf += dot(x[i00], x[i00]);
    }
    sumf = simd_sum(sumf);
    threadgroup_barrier(mem_flags::mem_threadgroup);
    if (tiisg == 0) shmem_f32[sgitg] = sumf;
    threadgroup_barrier(mem_flags::mem_threadgroup);
    sumf = simd_sum(shmem_f32[tiisg]);

    const float mean = sumf / args.ne00;
    const float scale = 1.0f / sqrt(mean + args.eps);
    if (tpitg.x == 0) dst_scale[i01] = scale;
}

struct ds4_metal_args_qkv_rms_norm {
    int32_t  q_n;
    int32_t  q_n4;
    int32_t  kv_n;
    int32_t  kv_n4;
    uint64_t q_row_stride;
    uint64_t kv_row_stride;
    float    eps;
};

// Normalizes DS4's q-lora row and KV row in one dispatch.  The two reductions
// deliberately mirror kernel_rms_norm_mul_f32_4: Q uses the full 256-thread
// row shape for 1024 floats, while KV only has work in the first 128 lanes for
// its 512 floats.  This keeps the q/kv normalization math aligned with the
// standalone kernels while removing one tiny launch from the attention setup.
kernel void kernel_dsv4_qkv_rms_norm_f32_4(
        constant ds4_metal_args_qkv_rms_norm & args,
        device const float4 * q_src,
        device const float4 * q_weight,
        device       float4 * q_dst,
        device const float4 * kv_src,
        device const float4 * kv_weight,
        device       float4 * kv_dst,
        threadgroup float * shmem_f32 [[threadgroup(0)]],
        uint3   tgpig[[threadgroup_position_in_grid]],
        ushort3 tpitg[[thread_position_in_threadgroup]],
        ushort  sgitg[[simdgroup_index_in_threadgroup]],
        ushort  tiisg[[thread_index_in_simdgroup]],
        ushort3 ntg[[threads_per_threadgroup]]) {
    if (sgitg == 0) {
        shmem_f32[tiisg] = 0.0f;
    }

    const uint row = tgpig.x;
    const bool kv_task = tgpig.y != 0;
    const int n = kv_task ? args.kv_n : args.q_n;
    const int n4 = kv_task ? args.kv_n4 : args.q_n4;
    const uint64_t row_stride4 = (kv_task ? args.kv_row_stride : args.q_row_stride) / sizeof(float4);

    device const float4 * x = kv_task ? kv_src + row * row_stride4 : q_src + row * row_stride4;
    device const float4 * w = kv_task ? kv_weight : q_weight;
    device       float4 * y = kv_task ? kv_dst + row * row_stride4 : q_dst + row * row_stride4;

    float sumf = 0.0f;
    for (int i = tpitg.x; i < n4; i += ntg.x) {
        const float4 v = x[i];
        sumf += dot(v, v);
    }
    sumf = simd_sum(sumf);

    threadgroup_barrier(mem_flags::mem_threadgroup);

    if (tiisg == 0) {
        shmem_f32[sgitg] = sumf;
    }

    threadgroup_barrier(mem_flags::mem_threadgroup);

    sumf = shmem_f32[tiisg];
    sumf = simd_sum(sumf);

#ifdef DS4_METAL_NORM_RSQRT_DISABLE
    // Match the formula used by kernel_rms_norm_fuse_impl above so both RMSNorm
    // entry points produce bit-identical scales. Hardware rsqrt() and 1.0f/sqrt()
    // can differ by ~1 ULP and that difference compounds across 43 layers.
    const float scale = 1.0f / sqrt(sumf / float(n) + args.eps);
#else
    const float scale = rsqrt(sumf / float(n) + args.eps);
#endif

    for (int i = tpitg.x; i < n4; i += ntg.x) {
        y[i] = (x[i] * scale) * w[i];
    }
}

// Decode-only triple fusion: the q/kv RMS norm, the KV RoPE tail, and the
// FP8/raw finalizer were three back-to-back dispatches on the same rows.
// The q threadgroup is byte-identical to kernel_dsv4_qkv_rms_norm_f32_4.
// The kv threadgroup continues with the shared affine-row RoPE helper (lane
// mapping preserved: r == lane on the first 64 lanes) and a verbatim copy of
// kernel_dsv4_kv_fp8_store_f32 with its work predicated to the first 64
// lanes (barriers stay uniform across the whole threadgroup).  Arithmetic,
// order and rounding are unchanged; gated and verified against
// full-vocabulary logits before promotion.
kernel void kernel_dsv4_qkv_rms_norm_kv_rope_fp8_store_f32(
        constant ds4_metal_args_qkv_rms_norm & args,
        constant ds4_metal_args_dsv4_rope_affine_pair & rope,
        constant ds4_metal_args_dsv4_kv_fp8_store & store,
        device const float4 * q_src,
        device const float4 * q_weight,
        device       float4 * q_dst,
        device const float4 * kv_src,
        device const float4 * kv_weight,
        device       float4 * kv_dst,
        device       float  * raw_cache,
        threadgroup float * shmem_f32 [[threadgroup(0)]],
        uint3   tgpig[[threadgroup_position_in_grid]],
        ushort3 tpitg[[thread_position_in_threadgroup]],
        ushort  sgitg[[simdgroup_index_in_threadgroup]],
        ushort  tiisg[[thread_index_in_simdgroup]],
        ushort3 ntg[[threads_per_threadgroup]]) {
    if (sgitg == 0) {
        shmem_f32[tiisg] = 0.0f;
    }

    const uint row = tgpig.x;
    const bool kv_task = tgpig.y != 0;
    const int n = kv_task ? args.kv_n : args.q_n;
    const int n4 = kv_task ? args.kv_n4 : args.q_n4;
    const uint64_t row_stride4 = (kv_task ? args.kv_row_stride : args.q_row_stride) / sizeof(float4);

    device const float4 * x = kv_task ? kv_src + row * row_stride4 : q_src + row * row_stride4;
    device const float4 * w = kv_task ? kv_weight : q_weight;
    device       float4 * y = kv_task ? kv_dst + row * row_stride4 : q_dst + row * row_stride4;

    float sumf = 0.0f;
    for (int i = tpitg.x; i < n4; i += ntg.x) {
        const float4 v = x[i];
        sumf += dot(v, v);
    }
    sumf = simd_sum(sumf);

    threadgroup_barrier(mem_flags::mem_threadgroup);

    if (tiisg == 0) {
        shmem_f32[sgitg] = sumf;
    }

    threadgroup_barrier(mem_flags::mem_threadgroup);

    sumf = shmem_f32[tiisg];
    sumf = simd_sum(sumf);

#ifdef DS4_METAL_NORM_RSQRT_DISABLE
    const float scale = 1.0f / sqrt(sumf / float(n) + args.eps);
#else
    const float scale = rsqrt(sumf / float(n) + args.eps);
#endif

    for (int i = tpitg.x; i < n4; i += ntg.x) {
        y[i] = (x[i] * scale) * w[i];
    }

    if (!kv_task) {
        return;
    }

    // KV RoPE tail in place, then the FP8/raw finalizer (verbatim bodies).
    threadgroup_barrier(mem_flags::mem_device_and_threadgroup);

    device char *kv_row = (device char *)(kv_dst + row * row_stride4);
    const int rope_n_nope = rope.head_dim - rope.n_dims;
    if (rope_n_nope < 0) {
        return;
    }
    ds4_rope_tail_pair_affine_row(rope,
                                  (device const char *)kv_row,
                                  kv_row,
                                  rope_n_nope,
                                  rope.pos0,
                                  tpitg.x,
                                  ntg.x);

    threadgroup_barrier(mem_flags::mem_device_and_threadgroup);

    const int head_dim = store.head_dim;
    const int n_rot = store.n_rot;
    const int n_nope = head_dim - n_rot;
    if (head_dim <= 0 || n_rot < 0 || n_nope < 0) {
        return;
    }
    const uint tid = tpitg.x;

    device float *kv = (device float *)kv_row;
    device float *raw = raw_cache + (int64_t)store.raw_row * head_dim;
    threadgroup float *scratch = shmem_f32 + 32;

    for (int off = 0; off < n_nope; off += 64) {
        float v = 0.0f;
        if (tid < 64u && off + (int)tid < n_nope) {
            v = kv[off + tid];
            scratch[tid] = abs(v);
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        for (uint stride = 32; stride > 0; stride >>= 1) {
            if (tid < stride) {
                scratch[tid] = max(scratch[tid], scratch[tid + stride]);
            }
            threadgroup_barrier(mem_flags::mem_threadgroup);
        }

        const float amax = max(scratch[0], 1.0e-4f);
        const float fp8_scale = exp2(ceil(log2(amax / 448.0f)));
        if (tid < 64u && off + (int)tid < n_nope) {
            const float q = dsv4_e4m3fn_dequant(clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale;
            kv[off + tid] = q;
#ifdef DS4_METAL_KV_RAW_F32
            raw[off + tid] = q;
#else
            raw[off + tid] = (float)((half)q);
#endif
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    if (tid < 64u) {
        for (int i = n_nope + tid; i < head_dim; i += 64) {
#ifdef DS4_METAL_KV_RAW_F32
            raw[i] = kv[i];
#else
            raw[i] = (float)((half)kv[i]);
#endif
        }
    }
}