#![allow(unsafe_code)] // Required for CUDA kernel launches.

//! Rust wrapper for MoE GPU kernels.
//! Loaded from cubin at runtime, like other cuda_pipeline heads (gpu_iqn_head, etc.).
//!
//! Per `feedback_no_htod_htoh_only_mapped_pinned.md`: all CPU<->GPU
//! communication uses mapped pinned memory. The test entry points
//! allocate MappedF32Buffers, write CPU data via host_ptr, launch the
//! kernel against dev_ptr, sync, and read GPU-written output via host_ptr.

use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};

use crate::MLError;
use super::mapped_pinned::MappedF32Buffer;

static MOE_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/moe_kernels.cubin"));

/// Adaptive load-balance λ controller (separate cubin — single-block,
/// single-thread cold-path producer matching kelly_cap_update precedent).
/// Reads ISV[MOE_GATE_ENTROPY_EMA_INDEX], writes ISV[MOE_LAMBDA_EFF_INDEX].
static MOE_LAMBDA_EFF_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/moe_lambda_eff_kernel.cubin"));

/// GPU-accelerated MoE head.
///
/// Owns the kernel function handles for the MoE mixture forward/backward pass,
/// load-balance loss computation, expert utilization EMA update, and softmax
/// forward/backward for the gate sub-network.
#[allow(missing_debug_implementations)]
pub struct GpuMoeHead {
    stream: Arc<CudaStream>,
    moe_mixture_forward: CudaFunction,
    moe_mixture_backward: CudaFunction,
    moe_dgate_reduce: CudaFunction,
    moe_load_balance_loss: CudaFunction,
    moe_load_balance_reduce: CudaFunction,
    moe_expert_util_ema_update: CudaFunction,
    /// Per-row softmax for the gate pre-activation → gate probabilities.
    moe_row_softmax: CudaFunction,
    /// Softmax Jacobian backward for the gate.
    moe_softmax_backward: CudaFunction,
    /// Adaptive load-balance λ controller — reads gate-entropy EMA from ISV,
    /// writes λ_eff = λ_floor + λ_max_extra × deficit into ISV[128].
    moe_lambda_eff_update: CudaFunction,
}

impl GpuMoeHead {
    /// Load MoE cubin and cache all kernel function handles.
    /// Mirrors the loader pattern in `GpuAuxHeads::new` and `GrnBlock::new`.
    pub fn new(stream: Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();
        let module = context
            .load_cubin(MOE_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("MoE cubin load: {e}")))?;
        let moe_mixture_forward = module
            .load_function("moe_mixture_forward")
            .map_err(|e| MLError::ModelError(format!("moe_mixture_forward load: {e}")))?;
        let moe_mixture_backward = module
            .load_function("moe_mixture_backward")
            .map_err(|e| MLError::ModelError(format!("moe_mixture_backward load: {e}")))?;
        let moe_dgate_reduce = module
            .load_function("moe_dgate_reduce")
            .map_err(|e| MLError::ModelError(format!("moe_dgate_reduce load: {e}")))?;
        let moe_load_balance_loss = module
            .load_function("moe_load_balance_loss")
            .map_err(|e| MLError::ModelError(format!("moe_load_balance_loss load: {e}")))?;
        let moe_load_balance_reduce = module
            .load_function("moe_load_balance_reduce")
            .map_err(|e| MLError::ModelError(format!("moe_load_balance_reduce load: {e}")))?;
        let moe_expert_util_ema_update = module
            .load_function("moe_expert_util_ema_update")
            .map_err(|e| MLError::ModelError(format!("moe_expert_util_ema_update load: {e}")))?;
        let moe_row_softmax = module
            .load_function("moe_row_softmax")
            .map_err(|e| MLError::ModelError(format!("moe_row_softmax load: {e}")))?;
        let moe_softmax_backward = module
            .load_function("moe_softmax_backward")
            .map_err(|e| MLError::ModelError(format!("moe_softmax_backward load: {e}")))?;
        // Adaptive λ controller — separate cubin (built from
        // moe_lambda_eff_kernel.cu by build.rs).
        let lambda_module = context
            .load_cubin(MOE_LAMBDA_EFF_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("moe_lambda_eff cubin load: {e}")))?;
        let moe_lambda_eff_update = lambda_module
            .load_function("moe_lambda_eff_update")
            .map_err(|e| MLError::ModelError(format!("moe_lambda_eff_update load: {e}")))?;
        Ok(Self {
            stream,
            moe_mixture_forward,
            moe_mixture_backward,
            moe_dgate_reduce,
            moe_load_balance_loss,
            moe_load_balance_reduce,
            moe_expert_util_ema_update,
            moe_row_softmax,
            moe_softmax_backward,
            moe_lambda_eff_update,
        })
    }

    // ══════════════════════════════════════════════════════════════════════
    // Production launch methods (operate on already-allocated device buffers)
    // ══════════════════════════════════════════════════════════════════════

    /// Launch `moe_row_softmax` in-place on already-allocated device pointers.
    /// Grid: (B+block-1)/block × 1 × 1. Each thread handles one row.
    pub(crate) fn launch_row_softmax(
        &self,
        in_pre_ptr: u64,
        out_soft_ptr: u64,
        b: usize,
        k: usize,
    ) -> Result<(), MLError> {
        let block: u32 = 256;
        let grid = ((b as u32) + block - 1) / block;
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        unsafe {
            self.stream
                .launch_builder(&self.moe_row_softmax)
                .arg(&in_pre_ptr)
                .arg(&out_soft_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .launch(LaunchConfig {
                    grid_dim: (grid, 1, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_row_softmax launch: {e}")))?;
        }
        Ok(())
    }

    /// Launch `moe_mixture_forward` on device pointers.
    /// h_s2_out[B,C] = sum_k gate[B,k] * expert_outputs[k,B,C].
    pub(crate) fn launch_mixture_forward(
        &self,
        expert_outputs_ptr: u64,
        gate_ptr: u64,
        h_s2_out_ptr: u64,
        b: usize,
        k: usize,
        c: usize,
    ) -> Result<(), MLError> {
        let total = (b * c) as u32;
        let block: u32 = 256;
        let grid = total.div_ceil(block);
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        let c_i32 = c as i32;
        unsafe {
            self.stream
                .launch_builder(&self.moe_mixture_forward)
                .arg(&expert_outputs_ptr)
                .arg(&gate_ptr)
                .arg(&h_s2_out_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&c_i32)
                .launch(LaunchConfig {
                    grid_dim: (grid, 1, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_mixture_forward launch: {e}")))?;
        }
        Ok(())
    }

    /// Launch `moe_load_balance_loss` + `moe_load_balance_reduce`.
    /// Writes the scalar total load-balance loss into `loss_total_ptr`.
    ///
    /// λ is read inside the kernel from `isv_signals[isv_lambda_eff_idx]`,
    /// where the adaptive controller (`moe_lambda_eff_update`) writes the
    /// effective load-balance weight each step. Per
    /// `feedback_isv_for_adaptive_bounds.md`, the load-balance push scales
    /// inversely with observed gate-entropy without any DtoH roundtrip.
    pub(crate) fn launch_load_balance_loss(
        &self,
        gate_ptr: u64,
        loss_per_k_ptr: u64,
        loss_total_ptr: u64,
        isv_signals_dev_ptr: u64,
        b: usize,
        k: usize,
        isv_lambda_eff_idx: usize,
    ) -> Result<(), MLError> {
        let block: u32 = 256;
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        let lambda_idx_i32 = isv_lambda_eff_idx as i32;
        unsafe {
            self.stream
                .launch_builder(&self.moe_load_balance_loss)
                .arg(&gate_ptr)
                .arg(&loss_per_k_ptr)
                .arg(&isv_signals_dev_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&lambda_idx_i32)
                .launch(LaunchConfig {
                    grid_dim: (k as u32, 1, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: block * 4,
                })
                .map_err(|e| MLError::ModelError(format!("moe_load_balance_loss launch: {e}")))?;

            self.stream
                .launch_builder(&self.moe_load_balance_reduce)
                .arg(&loss_per_k_ptr)
                .arg(&loss_total_ptr)
                .arg(&k_i32)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_load_balance_reduce launch: {e}")))?;
        }
        Ok(())
    }

    /// Launch `moe_lambda_eff_update` ISV producer on device pointers.
    ///
    /// Reads ISV[`isv_entropy_ema_idx`] (gate-entropy EMA from
    /// `moe_expert_util_ema_update`) and writes ISV[`isv_lambda_eff_idx`] =
    /// `lambda_floor + lambda_max_extra × deficit`, where `deficit ∈ [0, 1]`
    /// is the normalised entropy shortfall against
    /// `entropy_target_frac × ln(K)`. Single-block, single-thread cold-path
    /// kernel (matches kelly_cap_update / cql_alpha_seed_update shape).
    ///
    /// Per `pearl_blend_formulas_must_have_permanent_floor.md`,
    /// `lambda_floor` is a permanent minimum — λ_eff never drops below it.
    pub(crate) fn launch_lambda_eff_update(
        &self,
        isv_dev_ptr: u64,
        isv_lambda_eff_idx: usize,
        isv_entropy_ema_idx: usize,
        lambda_floor: f32,
        lambda_max_extra: f32,
        entropy_target_frac: f32,
    ) -> Result<(), MLError> {
        let lambda_idx_i32 = isv_lambda_eff_idx as i32;
        let entropy_idx_i32 = isv_entropy_ema_idx as i32;
        unsafe {
            self.stream
                .launch_builder(&self.moe_lambda_eff_update)
                .arg(&isv_dev_ptr)
                .arg(&lambda_idx_i32)
                .arg(&entropy_idx_i32)
                .arg(&lambda_floor)
                .arg(&lambda_max_extra)
                .arg(&entropy_target_frac)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_lambda_eff_update launch: {e}")))?;
        }
        Ok(())
    }

    /// Launch `moe_expert_util_ema_update` step-observation producer on
    /// device pointers. SP4 Layer A Task A13.2 retrofit (2026-05-01): writes
    /// per-expert col_mean to `scratch_buf[scratch_idx_util_base..+K)` and
    /// gate Shannon entropy to `scratch_buf[scratch_idx_entropy]`. The host
    /// applies Pearls A+D via `pearls_ad_update`.
    pub(crate) fn launch_expert_util_ema(
        &self,
        gate_ptr: u64,
        scratch_dev_ptr: u64,
        b: usize,
        k: usize,
        scratch_idx_util_base: usize,
        scratch_idx_entropy: usize,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        let util_base_i32 = scratch_idx_util_base as i32;
        let entropy_idx_i32 = scratch_idx_entropy as i32;
        unsafe {
            self.stream
                .launch_builder(&self.moe_expert_util_ema_update)
                .arg(&gate_ptr)
                .arg(&scratch_dev_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&util_base_i32)
                .arg(&entropy_idx_i32)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_expert_util_ema_update launch: {e}")))?;
        }
        Ok(())
    }

    /// Launch `moe_mixture_backward` (de_k) + `moe_dgate_reduce` (dg) on device pointers.
    pub(crate) fn launch_mixture_backward(
        &self,
        dh_s2_ptr: u64,
        gate_ptr: u64,
        expert_outputs_ptr: u64,
        de_k_ptr: u64,
        dg_ptr: u64,
        b: usize,
        k: usize,
        c: usize,
    ) -> Result<(), MLError> {
        let block: u32 = 256;
        let total_kbc = (k * b * c) as u32;
        let grid_kbc = total_kbc.div_ceil(block);
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        let c_i32 = c as i32;
        unsafe {
            self.stream
                .launch_builder(&self.moe_mixture_backward)
                .arg(&dh_s2_ptr)
                .arg(&gate_ptr)
                .arg(&de_k_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&c_i32)
                .launch(LaunchConfig {
                    grid_dim: (grid_kbc, 1, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_mixture_backward launch: {e}")))?;

            self.stream
                .launch_builder(&self.moe_dgate_reduce)
                .arg(&dh_s2_ptr)
                .arg(&expert_outputs_ptr)
                .arg(&dg_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&c_i32)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, k as u32, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: block * 4,
                })
                .map_err(|e| MLError::ModelError(format!("moe_dgate_reduce launch: {e}")))?;
        }
        Ok(())
    }

    /// Launch `moe_softmax_backward` on device pointers.
    pub(crate) fn launch_softmax_backward(
        &self,
        g_soft_ptr: u64,
        dg_ptr: u64,
        dg_pre_ptr: u64,
        b: usize,
        k: usize,
    ) -> Result<(), MLError> {
        let block: u32 = 256;
        let grid = ((b as u32) + block - 1) / block;
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        unsafe {
            self.stream
                .launch_builder(&self.moe_softmax_backward)
                .arg(&g_soft_ptr)
                .arg(&dg_ptr)
                .arg(&dg_pre_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .launch(LaunchConfig {
                    grid_dim: (grid, 1, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_softmax_backward launch: {e}")))?;
        }
        Ok(())
    }

    /// Test-only entry — stages CPU data via mapped pinned memory, runs the
    /// kernel, reads mapped pinned output. NO HtoD/HtoH per
    /// `feedback_no_htod_htoh_only_mapped_pinned.md`.
    ///
    /// Arguments:
    ///   `expert_outputs` — `[K, B, C]` row-major flat slice.
    ///   `gate`           — `[B, K]` row-major flat slice (rows must sum to 1).
    ///   `b`              — batch size.
    ///   `k`              — number of experts.
    ///   `c`              — hidden/channel dimension.
    ///
    /// Returns `h_s2` as a `[B, C]` flat `Vec<f32>`.
    pub fn test_mixture_forward(
        &self,
        expert_outputs: &[f32],
        gate: &[f32],
        b: usize,
        k: usize,
        c: usize,
    ) -> Result<Vec<f32>, MLError> {
        assert_eq!(
            expert_outputs.len(),
            k * b * c,
            "expert_outputs length mismatch: expected {} ({k}×{b}×{c}), got {}",
            k * b * c,
            expert_outputs.len()
        );
        assert_eq!(
            gate.len(),
            b * k,
            "gate length mismatch: expected {} ({b}×{k}), got {}",
            b * k,
            gate.len()
        );

        // Allocate mapped pinned buffers (CPU-and-GPU coherent).
        let expert_buf = unsafe { MappedF32Buffer::new(k * b * c) }
            .map_err(|e| MLError::ModelError(format!("expert buf alloc: {e}")))?;
        let gate_buf = unsafe { MappedF32Buffer::new(b * k) }
            .map_err(|e| MLError::ModelError(format!("gate buf alloc: {e}")))?;
        let h_s2_buf = unsafe { MappedF32Buffer::new(b * c) }
            .map_err(|e| MLError::ModelError(format!("h_s2 buf alloc: {e}")))?;

        // Direct host_ptr writes, no memcpy.
        expert_buf.write_from_slice(expert_outputs);
        gate_buf.write_from_slice(gate);

        let total = (b * c) as u32;
        let block: u32 = 256;
        let grid = total.div_ceil(block);

        let dev_expert = expert_buf.dev_ptr;
        let dev_gate = gate_buf.dev_ptr;
        let dev_h_s2 = h_s2_buf.dev_ptr;
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        let c_i32 = c as i32;

        unsafe {
            self.stream
                .launch_builder(&self.moe_mixture_forward)
                .arg(&dev_expert)
                .arg(&dev_gate)
                .arg(&dev_h_s2)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&c_i32)
                .launch(LaunchConfig {
                    grid_dim: (grid, 1, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_mixture_forward launch: {e}")))?;
        }
        self.stream
            .synchronize()
            .map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?;

        // Read GPU-written result via host_ptr (mapped pinned coherence).
        Ok(h_s2_buf.read_all())
    }

    /// Test-only entry for the backward pass.
    ///
    /// Launches `moe_mixture_backward` (de_k = g · dh_s2) and
    /// `moe_dgate_reduce` (dg = Σ_c e_k · dh_s2, shmem reduction over c).
    ///
    /// Returns `(de_k [K*B*C], dg [B*K])`.
    pub fn test_mixture_backward(
        &self,
        expert_outputs: &[f32],
        gate: &[f32],
        dh_s2: &[f32],
        b: usize,
        k: usize,
        c: usize,
    ) -> Result<(Vec<f32>, Vec<f32>), MLError> {
        let expert_buf = unsafe { MappedF32Buffer::new(k * b * c) }
            .map_err(|e| MLError::ModelError(format!("expert buf alloc: {e}")))?;
        let gate_buf = unsafe { MappedF32Buffer::new(b * k) }
            .map_err(|e| MLError::ModelError(format!("gate buf alloc: {e}")))?;
        let dh_s2_buf = unsafe { MappedF32Buffer::new(b * c) }
            .map_err(|e| MLError::ModelError(format!("dh_s2 buf alloc: {e}")))?;
        let de_k_buf = unsafe { MappedF32Buffer::new(k * b * c) }
            .map_err(|e| MLError::ModelError(format!("de_k buf alloc: {e}")))?;
        let dg_buf = unsafe { MappedF32Buffer::new(b * k) }
            .map_err(|e| MLError::ModelError(format!("dg buf alloc: {e}")))?;

        expert_buf.write_from_slice(expert_outputs);
        gate_buf.write_from_slice(gate);
        dh_s2_buf.write_from_slice(dh_s2);

        let block: u32 = 256;
        let total_kbc = (k * b * c) as u32;
        let grid_kbc = total_kbc.div_ceil(block);

        let dev_expert = expert_buf.dev_ptr;
        let dev_gate = gate_buf.dev_ptr;
        let dev_dh_s2 = dh_s2_buf.dev_ptr;
        let dev_de_k = de_k_buf.dev_ptr;
        let dev_dg = dg_buf.dev_ptr;
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        let c_i32 = c as i32;

        unsafe {
            self.stream
                .launch_builder(&self.moe_mixture_backward)
                .arg(&dev_dh_s2)
                .arg(&dev_gate)
                .arg(&dev_de_k)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&c_i32)
                .launch(LaunchConfig {
                    grid_dim: (grid_kbc, 1, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_mixture_backward launch: {e}")))?;

            self.stream
                .launch_builder(&self.moe_dgate_reduce)
                .arg(&dev_dh_s2)
                .arg(&dev_expert)
                .arg(&dev_dg)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&c_i32)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, k as u32, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: block * 4, // block * sizeof(f32)
                })
                .map_err(|e| MLError::ModelError(format!("moe_dgate_reduce launch: {e}")))?;
        }
        self.stream
            .synchronize()
            .map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?;

        Ok((de_k_buf.read_all(), dg_buf.read_all()))
    }

    /// Test-only entry for the load-balance loss.
    ///
    /// Computes `λ · K · Σ_k (mean_b g[b,k])²` using two kernels:
    /// `moe_load_balance_loss` (per-k contribution) and
    /// `moe_load_balance_reduce` (K-element scalar sum).
    ///
    /// `lambda` is staged into a synthetic ISV slot at index 0 of a small
    /// mapped pinned buffer; the kernel reads `isv_signals[0]` for λ. This
    /// mirrors the production contract where the consumer reads
    /// `ISV[MOE_LAMBDA_EFF_INDEX]` filled by the `moe_lambda_eff_update`
    /// adaptive controller.
    ///
    /// Returns the scalar total loss.
    pub fn test_load_balance_loss(
        &self,
        gate: &[f32],
        b: usize,
        k: usize,
        lambda: f32,
    ) -> Result<f32, MLError> {
        let gate_buf = unsafe { MappedF32Buffer::new(b * k) }
            .map_err(|e| MLError::ModelError(format!("gate buf alloc: {e}")))?;
        let loss_per_k_buf = unsafe { MappedF32Buffer::new(k) }
            .map_err(|e| MLError::ModelError(format!("loss_per_k buf alloc: {e}")))?;
        let loss_total_buf = unsafe { MappedF32Buffer::new(1) }
            .map_err(|e| MLError::ModelError(format!("loss_total buf alloc: {e}")))?;
        // Test-shim ISV: 1-element buffer with λ at index 0 so the kernel's
        // `isv_signals[isv_lambda_eff_idx]` read returns the requested value.
        let isv_buf = unsafe { MappedF32Buffer::new(1) }
            .map_err(|e| MLError::ModelError(format!("isv buf alloc: {e}")))?;

        gate_buf.write_from_slice(gate);
        isv_buf.write_from_slice(&[lambda]);

        let block: u32 = 256;
        let dev_gate = gate_buf.dev_ptr;
        let dev_loss_per_k = loss_per_k_buf.dev_ptr;
        let dev_loss_total = loss_total_buf.dev_ptr;
        let dev_isv = isv_buf.dev_ptr;
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        let isv_lambda_idx_i32: i32 = 0;

        unsafe {
            self.stream
                .launch_builder(&self.moe_load_balance_loss)
                .arg(&dev_gate)
                .arg(&dev_loss_per_k)
                .arg(&dev_isv)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&isv_lambda_idx_i32)
                .launch(LaunchConfig {
                    grid_dim: (k as u32, 1, 1),
                    block_dim: (block, 1, 1),
                    shared_mem_bytes: block * 4, // block * sizeof(f32)
                })
                .map_err(|e| MLError::ModelError(format!("moe_load_balance_loss launch: {e}")))?;

            self.stream
                .launch_builder(&self.moe_load_balance_reduce)
                .arg(&dev_loss_per_k)
                .arg(&dev_loss_total)
                .arg(&k_i32)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_load_balance_reduce launch: {e}")))?;
        }
        self.stream
            .synchronize()
            .map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?;

        Ok(loss_total_buf.read_all()[0])
    }

    /// Test-only entry for the ISV EMA update kernel.
    ///
    /// Writes 8 per-expert utilization EMAs + 1 gate-entropy EMA into the
    /// provided ISV buffer at `isv_util_base..isv_util_base+K` and
    /// `isv_entropy_index`, with decay `alpha`.
    ///
    /// GPU-resident per `pearl_cold_path_no_exception_to_gpu_drives.md`.
    /// Returns the full updated ISV as a `Vec<f32>`.
    pub fn test_expert_util_ema_update(
        &self,
        isv: &[f32],
        gate: &[f32],
        b: usize,
        k: usize,
        alpha: f32,
        isv_util_base: usize,
        isv_entropy_index: usize,
    ) -> Result<Vec<f32>, MLError> {
        let isv_buf = unsafe { MappedF32Buffer::new(isv.len()) }
            .map_err(|e| MLError::ModelError(format!("isv buf alloc: {e}")))?;
        let gate_buf = unsafe { MappedF32Buffer::new(b * k) }
            .map_err(|e| MLError::ModelError(format!("gate buf alloc: {e}")))?;

        isv_buf.write_from_slice(isv);
        gate_buf.write_from_slice(gate);

        let dev_gate = gate_buf.dev_ptr;
        let dev_isv = isv_buf.dev_ptr;
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        let util_base_i32 = isv_util_base as i32;
        let entropy_idx_i32 = isv_entropy_index as i32;

        unsafe {
            self.stream
                .launch_builder(&self.moe_expert_util_ema_update)
                .arg(&dev_gate)
                .arg(&dev_isv)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&util_base_i32)
                .arg(&entropy_idx_i32)
                .arg(&alpha)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_expert_util_ema_update launch: {e}")))?;
        }
        self.stream
            .synchronize()
            .map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?;

        Ok(isv_buf.read_all())
    }

    /// Test-only entry for the adaptive λ controller kernel.
    ///
    /// Stages the provided `isv` slice into mapped pinned memory, runs
    /// `moe_lambda_eff_update`, and returns the full updated ISV. The
    /// controller reads `isv[isv_entropy_ema_idx]` and writes
    /// `isv[isv_lambda_eff_idx] = lambda_floor + lambda_max_extra × deficit`,
    /// where `deficit = clamp((target − ent_ema)/target, 0, 1)` and
    /// `target = entropy_target_frac × ln(K=8)`.
    ///
    /// GPU-resident per `pearl_cold_path_no_exception_to_gpu_drives.md`.
    pub fn test_lambda_eff_update(
        &self,
        isv: &[f32],
        isv_lambda_eff_idx: usize,
        isv_entropy_ema_idx: usize,
        lambda_floor: f32,
        lambda_max_extra: f32,
        entropy_target_frac: f32,
    ) -> Result<Vec<f32>, MLError> {
        let isv_buf = unsafe { MappedF32Buffer::new(isv.len()) }
            .map_err(|e| MLError::ModelError(format!("isv buf alloc: {e}")))?;
        isv_buf.write_from_slice(isv);

        let dev_isv = isv_buf.dev_ptr;
        let lambda_idx_i32 = isv_lambda_eff_idx as i32;
        let entropy_idx_i32 = isv_entropy_ema_idx as i32;

        unsafe {
            self.stream
                .launch_builder(&self.moe_lambda_eff_update)
                .arg(&dev_isv)
                .arg(&lambda_idx_i32)
                .arg(&entropy_idx_i32)
                .arg(&lambda_floor)
                .arg(&lambda_max_extra)
                .arg(&entropy_target_frac)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("moe_lambda_eff_update launch: {e}")))?;
        }
        self.stream
            .synchronize()
            .map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?;

        Ok(isv_buf.read_all())
    }
}
