mlx-native 0.21.0

//! GGML block-format expert-routed (MoE) quantized matrix-vector multiply dispatch.
//!
//! Encodes a GPU compute command that performs, for each (token, expert-slot):
//!   expert_id = ids[token * top_k + slot]
//!   output[token*top_k + slot][col] = sum_k(dequant(weight[expert_id][col][k]) * input[token][k])
//!
//! This is the _id variant of quantized_matmul_ggml: same GGML block dequantization
//! but with per-token expert selection via an ids buffer, enabling fused MoE dispatch.
//!
//! Derived from candle-metal-kernels (Apache-2.0) kernel_mul_mv_id template
//! and mlx-native's quantized_matmul_ggml kernels.

use crate::buffer::MlxBuffer;
use crate::device::MlxDevice;
use crate::dtypes::DType;
use crate::encoder::{as_bytes, CapturedOpKind, CommandEncoder, DispatchRecord, KernelArg};
use crate::env_flags::{cached_env_default_true, cached_env_eq_one};
use crate::ggml_capability::{
    compact_max_tiles, expert_mm_schedule_scratch_bytes, ggml_expert_bytes, plan_expert_auto_route,
    tensor_mm_auto_selected, ExpertAutoPlan, ExpertMmSchedule, GgmlRoutingPolicy,
    GgmlTensorMmPreference,
};
use crate::ggml_routing_policy::ggml_routing_policy_for_registry;
use std::sync::atomic::AtomicI8;

// ADR-029: cached hot-path env-flag gates for dispatch_id_mv.
// Uncached env::var costs ~70 ns/call (H-N microbench); cached AtomicI8 ~2 ns.
// dispatch_id_mv is called ~150 times/token; 2 env reads × 150 = 300 × 68 ns
// = 20 µs/tok ≈ 0.2% wall just from these 2 sites.
static CACHED_Q6K_ID_MV_NR2: AtomicI8 = AtomicI8::new(-1);
static CACHED_Q8_0_ID_MV_NR2: AtomicI8 = AtomicI8::new(-1);
use crate::error::{MlxError, Result};
use crate::kernel_registry::KernelRegistry;
use crate::ops::dense_mm_capability::is_unavailable_tensor_header;
use crate::ops::quantized_matmul_ggml::GgmlType;

fn checked_byte_extent(label: &str, factors: &[usize]) -> Result<usize> {
    factors.iter().try_fold(1usize, |total, factor| {
        total.checked_mul(*factor).ok_or_else(|| {
            MlxError::InvalidArgument(format!("quantized_matmul_id_ggml: {label} size overflow"))
        })
    })
}

// ---- GPU params struct ----

/// GPU-side params struct — must match the Metal shader's `GgmlMatvecIdParams`.
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct GgmlMatvecIdGpuParams {
    ne00: i64,          // K
    ne01: i64,          // N
    ne02: i64,          // number of experts (on-device ID bound)
    ne10: i64,          // K
    ne12: i64,          // 1 (unused)
    ne0: i64,           // N (output stride)
    ne1: i64,           // total output rows = n_tokens * top_k
    r2: u32,            // 1
    r3: u32,            // 1
    top_k: u32,         // experts per token
    n_tokens: u32,      // number of input tokens
    expert_stride: i64, // bytes between expert weight slices
}

// ---- Public types ----

/// Parameters describing the expert-routed GGML quantized matmul dimensions.
#[derive(Debug, Clone, Copy)]
pub struct GgmlQuantizedMatmulIdParams {
    /// Number of input tokens.
    pub n_tokens: u32,
    /// Number of experts each token is routed to (top-k).
    pub top_k: u32,
    /// Number of output columns per expert (weight rows).
    pub n: u32,
    /// Input dimension (weight cols before quantization).
    /// Must be divisible by the GGML block QK value.
    pub k: u32,
    /// Total number of experts in the stacked weight buffer.
    pub n_experts: u32,
    /// Byte stride between expert weight slices in the stacked buffer.
    pub expert_stride: u64,
    /// GGML quantization type.
    pub ggml_type: GgmlType,
}

impl GgmlType {
    /// Metal kernel function name for the mat-vec `_id` variant.
    fn id_kernel_name(self) -> &'static str {
        match self {
            GgmlType::Q4_0 => "kernel_mul_mv_id_q4_0_f32",
            GgmlType::Q5_0 => "kernel_mul_mv_id_q5_0_f32",
            GgmlType::Q8_0 => "kernel_mul_mv_id_q8_0_f32",
            GgmlType::Q2_K => "kernel_mul_mv_id_q2_K_f32",
            GgmlType::Q3_K => "kernel_mul_mv_id_q3_K_f32",
            // ADR-013 P7 — Q4_K mv_id peer-pattern kernel for
            // dwq46/dwq48 MoE expert weights.
            GgmlType::Q4_K => "kernel_mul_mv_id_q4_K_f32",
            GgmlType::Q5_K => "kernel_mul_mv_id_q5_K_f32",
            GgmlType::Q6_K => "kernel_mul_mv_id_q6_K_f32",
            // ADR-022 Phase 1 —Q5_1 / IQ4_NL mv_id ports.
            GgmlType::Q5_1 => "kernel_mul_mv_id_q5_1_f32",
            GgmlType::IQ4_NL => "kernel_mul_mv_id_iq4_nl_f32",
            GgmlType::F32 | GgmlType::F16 | GgmlType::BF16 | GgmlType::I16 | GgmlType::I32 => {
                "unsupported"
            }
            // ADR-033 §Pi Task #16 SHIPPED 2026-05-22 — mirrors IQ4_NL
            // _id geometry (N_SIMDGROUP=2, N_DST=4, (8, 8, 8) launch).
            GgmlType::IQ4_XS => "kernel_mul_mv_id_iq4_xs_f32",
            // ADR-066 N-IQ2XXS — port of the peer's IQ2_XXS mv layout
            // (N_R0 = 4, N_SG = 2, grid/sign tables in threadgroup memory).
            GgmlType::IQ2_XXS => "kernel_mul_mv_id_iq2_xxs_f32",
        }
    }

    /// Metal kernel function name for the mat-mat `_id` variant (ADR-011
    /// Phase 3 Wave P3a peer port of `kernel_mul_mm_id_<q>_f32`).
    fn id_mm_kernel_name(self) -> &'static str {
        match self {
            GgmlType::Q4_0 => "kernel_mul_mm_id_q4_0_f32",
            GgmlType::Q5_0 => "kernel_mul_mm_id_q5_0_f32",
            GgmlType::Q8_0 => "kernel_mul_mm_id_q8_0_f32",
            GgmlType::Q2_K => "kernel_mul_mm_id_q2_K_f32",
            GgmlType::Q3_K => "kernel_mul_mm_id_q3_K_f32",
            // ADR-022 Phase 2 — Q5_K mm_id ported.
            GgmlType::Q5_K => "kernel_mul_mm_id_q5_K_f32",
            GgmlType::Q6_K => "kernel_mul_mm_id_q6_K_f32",
            // ADR-013 P16 — Q4_K mm_id peer port of
            // `kernel_mul_mm_id_q4_K_f32`.
            GgmlType::Q4_K => "kernel_mul_mm_id_q4_K_f32",
            // ADR-022 Phase 1 —Q5_1 / IQ4_NL mm_id ported.
            GgmlType::Q5_1 => "kernel_mul_mm_id_q5_1_f32",
            GgmlType::IQ4_NL => "kernel_mul_mm_id_iq4_nl_f32",
            GgmlType::F32 | GgmlType::F16 | GgmlType::BF16 | GgmlType::I16 | GgmlType::I32 => {
                "unsupported"
            }
            // ADR-033 §Pi Task #20 — IQ4_XS mm_id ported (simdgroup MMA
            // path). Tensor-API variant is not yet ported; tests fall
            // back to the simdgroup path via TENSOR_MM_ID_AVAILABLE probe.
            GgmlType::IQ4_XS => "kernel_mul_mm_id_iq4_xs_f32",
            // ADR-066 N-IQ2XXS — simdgroup MMA mm_id.
            GgmlType::IQ2_XXS => "kernel_mul_mm_id_iq2_xxs_f32",
        }
    }

    /// Tensor-API variant of the mm_id kernel (ADR-011 Phase 3 Wave
    /// P3b-tensor).
    fn id_mm_tensor_kernel_name(self) -> &'static str {
        match self {
            GgmlType::Q4_0 => "kernel_mul_mm_id_q4_0_tensor_f32",
            GgmlType::Q5_0 => "kernel_mul_mm_id_q5_0_tensor_f32",
            GgmlType::Q8_0 => "kernel_mul_mm_id_q8_0_tensor_f32",
            GgmlType::Q2_K => "kernel_mul_mm_id_q2_K_tensor_f32",
            GgmlType::Q3_K => "kernel_mul_mm_id_q3_K_tensor_f32",
            // ADR-022 Phase 2 — Q5_K mm_id_tensor ported.
            GgmlType::Q5_K => "kernel_mul_mm_id_q5_K_tensor_f32",
            GgmlType::Q6_K => "kernel_mul_mm_id_q6_K_tensor_f32",
            // ADR-013 P16 — Q4_K tensor-API mm_id ported.
            GgmlType::Q4_K => "kernel_mul_mm_id_q4_K_tensor_f32",
            // ADR-022 Phase 1 —Q5_1 / IQ4_NL tensor-API mm_id ported.
            GgmlType::Q5_1 => "kernel_mul_mm_id_q5_1_tensor_f32",
            GgmlType::IQ4_NL => "kernel_mul_mm_id_iq4_nl_tensor_f32",
            GgmlType::F32 | GgmlType::F16 | GgmlType::BF16 | GgmlType::I16 | GgmlType::I32 => {
                "unsupported"
            }
            // ADR-033 §Pi Task #20 tensor-API — IQ4_XS tensor-API mm_id
            // SHIPPED 2026-05-22 to close the prefill perf gap vs the
            // reference implementation.
            GgmlType::IQ4_XS => "kernel_mul_mm_id_iq4_xs_tensor_f32",
            // ADR-066 N-IQ2XXS — tensor-API mm_id.
            GgmlType::IQ2_XXS => "kernel_mul_mm_id_iq2_xxs_tensor_f32",
        }
    }

    /// Compact-schedule (`_c`) simdgroup mm_id kernel (ADR-063 P2).
    fn id_mm_compact_kernel_name(self) -> &'static str {
        match self {
            GgmlType::Q4_0 => "kernel_mul_mm_id_q4_0_f32_c",
            GgmlType::Q5_0 => "kernel_mul_mm_id_q5_0_f32_c",
            GgmlType::Q8_0 => "kernel_mul_mm_id_q8_0_f32_c",
            GgmlType::Q2_K => "kernel_mul_mm_id_q2_K_f32_c",
            GgmlType::Q3_K => "kernel_mul_mm_id_q3_K_f32_c",
            GgmlType::Q4_K => "kernel_mul_mm_id_q4_K_f32_c",
            GgmlType::Q5_K => "kernel_mul_mm_id_q5_K_f32_c",
            GgmlType::Q6_K => "kernel_mul_mm_id_q6_K_f32_c",
            GgmlType::Q5_1 => "kernel_mul_mm_id_q5_1_f32_c",
            GgmlType::IQ4_NL => "kernel_mul_mm_id_iq4_nl_f32_c",
            GgmlType::IQ4_XS => "kernel_mul_mm_id_iq4_xs_f32_c",
            GgmlType::IQ2_XXS => "kernel_mul_mm_id_iq2_xxs_f32_c",
            GgmlType::F32 | GgmlType::F16 | GgmlType::BF16 | GgmlType::I16 | GgmlType::I32 => {
                "unsupported"
            }
        }
    }

    /// Compact-schedule (`_c`) tensor-API mm_id kernel (ADR-063 P2).
    fn id_mm_tensor_compact_kernel_name(self) -> &'static str {
        match self {
            GgmlType::Q4_0 => "kernel_mul_mm_id_q4_0_tensor_f32_c",
            GgmlType::Q5_0 => "kernel_mul_mm_id_q5_0_tensor_f32_c",
            GgmlType::Q8_0 => "kernel_mul_mm_id_q8_0_tensor_f32_c",
            GgmlType::Q2_K => "kernel_mul_mm_id_q2_K_tensor_f32_c",
            GgmlType::Q3_K => "kernel_mul_mm_id_q3_K_tensor_f32_c",
            GgmlType::Q4_K => "kernel_mul_mm_id_q4_K_tensor_f32_c",
            GgmlType::Q5_K => "kernel_mul_mm_id_q5_K_tensor_f32_c",
            GgmlType::Q6_K => "kernel_mul_mm_id_q6_K_tensor_f32_c",
            GgmlType::Q5_1 => "kernel_mul_mm_id_q5_1_tensor_f32_c",
            GgmlType::IQ4_NL => "kernel_mul_mm_id_iq4_nl_tensor_f32_c",
            GgmlType::IQ4_XS => "kernel_mul_mm_id_iq4_xs_tensor_f32_c",
            GgmlType::IQ2_XXS => "kernel_mul_mm_id_iq2_xxs_tensor_f32_c",
            GgmlType::F32 | GgmlType::F16 | GgmlType::BF16 | GgmlType::I16 | GgmlType::I32 => {
                "unsupported"
            }
        }
    }
}

#[inline]
fn has_mm_id_map0(top_k: u32) -> bool {
    // Exact template set used by the production models currently supported:
    // top-k 6 is DeepSeek-V4; 1 and 8 retain the existing Gemma/Qwen paths.
    matches!(top_k, 1 | 6 | 8)
}

fn required_expert_weight_bytes(
    ggml_type: GgmlType,
    n_experts: u32,
    n: u32,
    k: u32,
    expert_stride: u64,
) -> Result<usize> {
    if expert_stride > i64::MAX as u64 {
        return Err(MlxError::InvalidArgument(
            "expert stride exceeds the signed Metal kernel ABI".into(),
        ));
    }
    let bytes = ggml_expert_bytes(ggml_type, n_experts, n, k, expert_stride)?;
    usize::try_from(bytes)
        .map_err(|_| MlxError::InvalidArgument("expert GGUF bytes exceed usize".into()))
}

fn validate_native_expert_dtypes(
    operation: &str,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
) -> Result<()> {
    if input.dtype() != DType::F32
        || weight.dtype() != DType::U8
        || ids.dtype() != DType::U32
        || output.dtype() != DType::F32
    {
        return Err(MlxError::InvalidArgument(format!(
            "{operation} requires F32 input, native U8 GGUF blocks, U32 expert IDs, and F32 output; got {:?}/{:?}/{:?}/{:?}",
            input.dtype(),
            weight.dtype(),
            ids.dtype(),
            output.dtype(),
        )));
    }
    Ok(())
}

fn validate_id_mm_scratch_dtypes(
    operation: &str,
    htpe: &MlxBuffer,
    hids: &MlxBuffer,
) -> Result<()> {
    if htpe.dtype() != DType::U32 || hids.dtype() != DType::U32 {
        return Err(MlxError::InvalidArgument(format!(
            "{operation} requires U32 htpe/hids scratch, got {:?}/{:?}",
            htpe.dtype(),
            hids.dtype(),
        )));
    }
    Ok(())
}

fn probe_tensor_mm_id(registry: &mut KernelRegistry, device: &MlxDevice) -> Result<bool> {
    let probe = registry.probe_optional_pipeline(
        "kernel_mul_mm_id_q4_0_tensor_f32",
        device.metal_device(),
        device.registry_id(),
        is_unavailable_tensor_header,
    )?;
    if probe.newly_probed && std::env::var("MLX_LOG_TENSOR_PROBE").is_ok() {
        eprintln!(
            "[mlx-native] tensor_mm_id probe: {}",
            if probe.available {
                "OK (using tensor variant for MoE)"
            } else {
                "FAILED (falling back to simdgroup MMA)"
            }
        );
    }
    Ok(probe.available)
}

pub(crate) fn expert_routing_policy_from_environment() -> GgmlRoutingPolicy {
    GgmlRoutingPolicy {
        expert_mm_threshold: mm_id_routing_threshold(),
        expert_q6k_mv_nr2: cached_env_default_true(&CACHED_Q6K_ID_MV_NR2, "HF2Q_Q6K_ID_MV_NR2"),
        expert_q8_0_mv_nr2: cached_env_eq_one(&CACHED_Q8_0_ID_MV_NR2, "HF2Q_Q8_0_ID_MV_NR2"),
        expert_tensor_mm: if std::env::var("HF2Q_DISABLE_TENSOR_MM_ID").is_ok() {
            GgmlTensorMmPreference::ForceSimd
        } else {
            GgmlTensorMmPreference::AutoProbe
        },
        q5_0_tensor_mm: std::env::var("HF2Q_Q5_0_TENSOR_MM").as_deref() == Ok("1"),
        expert_mm_schedule: mm_id_schedule_from_environment(),
        ..GgmlRoutingPolicy::default()
    }
}

/// `HF2Q_MM_ID_SCHEDULE`: `grid` selects [`ExpertMmSchedule::UniformGrid`]
/// (diagnostic and A/B arm), `compact` selects
/// [`ExpertMmSchedule::CompactTiles`]; unset or empty keeps the default.
/// Any other value warns once and keeps the default. Read once per process.
fn mm_id_schedule_from_environment() -> ExpertMmSchedule {
    static SCHEDULE: std::sync::LazyLock<ExpertMmSchedule> = std::sync::LazyLock::new(|| {
        let value = std::env::var("HF2Q_MM_ID_SCHEDULE").ok();
        let (schedule, recognized) = parse_mm_id_schedule(value.as_deref());
        if !recognized {
            eprintln!(
                "[mlx-native] HF2Q_MM_ID_SCHEDULE={:?} is neither `grid` nor `compact`; keeping the default {:?} schedule",
                value.unwrap_or_default(),
                schedule,
            );
        }
        schedule
    });
    *SCHEDULE
}

/// Parse an `HF2Q_MM_ID_SCHEDULE` value; the flag is `false` for a
/// non-empty value that is not a schedule name.
fn parse_mm_id_schedule(value: Option<&str>) -> (ExpertMmSchedule, bool) {
    let default = GgmlRoutingPolicy::default().expert_mm_schedule;
    match value.map(str::trim) {
        None | Some("") => (default, true),
        Some("grid") => (ExpertMmSchedule::UniformGrid, true),
        Some("compact") => (ExpertMmSchedule::CompactTiles, true),
        Some(_) => (default, false),
    }
}

/// Encode an expert-routed GGML quantized matrix-vector multiply.
///
/// Weight buffer contains raw GGML blocks stacked as `[n_experts, N, packed_K]`.
/// Input is f32 `[n_tokens, K]`, output is f32 `[n_tokens * top_k, N]`.
/// The `ids` buffer `[n_tokens * top_k]` u32 selects which expert to use for
/// each (token, slot) pair.
///
/// # Arguments
///
/// * `encoder`  -- Command encoder to record the dispatch into.
/// * `registry` -- Kernel registry (compiles shader on first call).
/// * `device`   -- Metal device.
/// * `input`    -- f32 input buffer, shape `[n_tokens, K]`.
/// * `weight`   -- Stacked GGML block weight buffer, `[n_experts, N, packed_K]`.
/// * `ids`      -- u32 expert index buffer, shape `[n_tokens * top_k]`.
/// * `output`   -- f32 output buffer, shape `[n_tokens * top_k, N]`.
/// * `params`   -- Dimensions and quantization parameters.
///
/// # Errors
///
/// Returns `MlxError::InvalidArgument` if:
/// - K is not divisible by the GGML block QK value
/// - Buffer sizes don't match expected dimensions
/// - Any dimension is zero
///
/// Expert IDs may be produced by an earlier GPU operation in the same
/// `CommandEncoder`, so this API does not synchronize to inspect them on the
/// CPU. The consumer kernels bounds-check IDs on-device. Matrix-matrix routes
/// also enforce distinct IDs per token. Invalid routing metadata performs no
/// expert-weight read and writes a non-finite poison sentinel: the affected
/// row for matvec, or the entire operation output for matrix-matrix. A graph
/// consumer must reject that sentinel at its existing result-readback boundary;
/// this encode-only API cannot turn device-produced metadata into a host `Err`.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlQuantizedMatmulIdParams,
) -> Result<()> {
    let routing = ggml_routing_policy_for_registry(registry);
    quantized_matmul_id_ggml_with_policy(
        encoder, registry, device, input, weight, ids, output, params, &routing,
    )
}

/// Expert auto-dispatch under an explicit receipt-bindable routing policy.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_with_policy(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlQuantizedMatmulIdParams,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    quantized_matmul_id_ggml_impl(
        encoder, registry, device, input, weight, ids, output, params, false, routing,
    )
}

/// ADR-040 Phase F `iter-F-moe-mvid` — byte-identity-required variant: always
/// routes to the per-token `mv_id` kernel, bypassing the `n_tokens > 32`
/// `mm_id` grouped route. Decode continuous-batching (the gemma4 `[N,hidden]`
/// batched body) at N≥5 has `n_tokens = N*top_k > 32` on the MoE down
/// projection; the `mm_id` grouped reduction is NOT bit-identical to the
/// serial per-token `mv_id` path, so a batched decode forward that must stay
/// byte-identical to N serial decodes calls this instead. `mm_id` is a prefill
/// throughput optimization and a measured regression at N≤8, so forcing `mv_id`
/// costs nothing at decode width.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_mv(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlQuantizedMatmulIdParams,
) -> Result<()> {
    let routing = ggml_routing_policy_for_registry(registry);
    quantized_matmul_id_ggml_mv_with_policy(
        encoder, registry, device, input, weight, ids, output, params, &routing,
    )
}

/// Forced expert matvec under an explicit routing policy.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_mv_with_policy(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlQuantizedMatmulIdParams,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    quantized_matmul_id_ggml_impl(
        encoder, registry, device, input, weight, ids, output, params, true, routing,
    )
}

/// Q4-GEMMA-DECODE F2: expert-grouped Q4_K gate_up with the
/// `moe_swiglu_batch` GELU-tanh(gate) * up epilogue, at the frozen geometry
/// of 4 gate/up row pairs per SIMD group (8 output columns per threadgroup).
///
/// The weight is the stacked `[n_experts, n, K]` gate_up tensor whose rows
/// `0..n/2` are gate and `n/2..n` are up; `output` is `[n_tokens * top_k,
/// n/2]` f32. The result is bitwise equal to `kernel_mul_mv_id_q4_K_f32`
/// into `[pairs, n]` followed by `moe_swiglu_batch_encode(.., n/2, pairs)`,
/// including out-of-range ids (NaN poison mapped through the epilogue).
/// Callers choose the decode widths that take this route.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_mv_grouped_geglu(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlQuantizedMatmulIdParams,
) -> Result<()> {
    /// Gate/up row pairs per SIMD group; the kernel runs two SIMD groups.
    const ROWS_PER_SIMDGROUP: usize = 4;
    if params.ggml_type != GgmlType::Q4_K || params.n % 2 != 0 {
        return Err(MlxError::InvalidArgument(format!(
            "grouped GEGLU mv_id requires Q4_K and an even gate_up width; got {:?} at N = {}",
            params.ggml_type, params.n
        )));
    }
    let n_ff = params.n / 2;
    validate_id_call(input, weight, ids, output, params, n_ff)?;

    let pipeline = registry.get_pipeline(
        "kernel_mul_mv_id_q4_K_f32_grouped_geglu_rt4",
        device.metal_device(),
    )?;
    let gpu_params = GgmlMatvecIdGpuParams {
        ne0: i64::from(n_ff),
        ..mv_id_gpu_params(params)
    };
    let total_rows = u64::from(params.n_tokens) * u64::from(params.top_k);
    encoder.encode_threadgroups_with_args(
        pipeline,
        &[
            (0, KernelArg::Buffer(weight)),
            (1, KernelArg::Buffer(input)),
            (2, KernelArg::Buffer(output)),
            (3, KernelArg::Buffer(ids)),
            (4, KernelArg::Bytes(as_bytes(&gpu_params))),
        ],
        metal::MTLSize::new(
            div_ceil(n_ff as usize, 2 * ROWS_PER_SIMDGROUP) as u64,
            total_rows,
            1,
        ),
        metal::MTLSize::new(2, 32, 1),
    );
    Ok(())
}

/// ADR-063 Q4-QWEN36 F2b: forced Q4_K expert mat-vec at the peer's layout
/// (`kernel_mul_mv_id_q4_K_f32_nr2`: two SIMD groups, two rows per 8-lane
/// block group, eight lanes per super-block, uint16 quant loads). Rows of
/// one or two super-blocks (K = 256, 512) fold the lane groups the peer
/// leaves idle onto further row pairs (4 / nb row groups per SIMD group).
///
/// One threadgroup computes one `(pair, column tile)` from that pair's
/// expert slice and token row, so each pair's output is bitwise independent
/// of `n_tokens * top_k` and of the other pairs' ids. Not bitwise equal to
/// `kernel_mul_mv_id_q4_K_f32` (different lane split and accumulation
/// order). Out-of-range ids write NaN like the stock kernel. Never chosen by
/// the default routing; callers that need one kernel at every width opt in.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_mv_q4_k_nr2(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlQuantizedMatmulIdParams,
) -> Result<()> {
    // Output columns per threadgroup: 2 SIMD groups x 2 rows x the kernel's
    // row-group fold (4 / nb for nb = 1, 2 super-blocks per row, else 1).
    let fold: usize = match params.k / 256 {
        1 => 4,
        2 => 2,
        _ => 1,
    };
    let columns_per_threadgroup = 4 * fold;
    if params.ggml_type != GgmlType::Q4_K {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml_mv_q4_k_nr2 requires Q4_K, got {:?}",
            params.ggml_type
        )));
    }
    validate_id_call(input, weight, ids, output, params, params.n)?;

    let pipeline = registry.get_pipeline("kernel_mul_mv_id_q4_K_f32_nr2", device.metal_device())?;
    let gpu_params = mv_id_gpu_params(params);
    let total_rows = u64::from(params.n_tokens) * u64::from(params.top_k);
    encoder.encode_threadgroups_with_args(
        pipeline,
        &[
            (0, KernelArg::Buffer(weight)),
            (1, KernelArg::Buffer(input)),
            (2, KernelArg::Buffer(output)),
            (3, KernelArg::Buffer(ids)),
            (4, KernelArg::Bytes(as_bytes(&gpu_params))),
        ],
        metal::MTLSize::new(
            div_ceil(params.n as usize, columns_per_threadgroup) as u64,
            total_rows,
            1,
        ),
        metal::MTLSize::new(32, 2, 1),
    );
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn quantized_matmul_id_ggml_impl(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlQuantizedMatmulIdParams,
    force_mv: bool,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    validate_id_call(input, weight, ids, output, params, params.n)?;

    // ADR-011 Phase 3 — route on n_tokens threshold.
    //
    // At prefill with n_tokens > 8, dispatch the two-stage mm_id kernels
    // (map0 + mm).  Each expert's weight tile is staged once to
    // threadgroup shmem per 32-row block of that expert's routed tokens
    // — identical to the dense mm dispatcher's win but per-expert.
    //
    // P3b-tensor.2 — extended to top_k=1 (Gemma 4's MoE down call).
    // Without this the down call's 19,640-row matmul falls back to mv_id
    // and re-reads each expert's weights once per row — ~50% of prefill
    // wall time burnt on weight re-reads.  The map0 template set is
    // {1, 6, 8}; other top_k values still fall back to mv_id.
    //
    // Falls back to the mv_id path for:
    //   * decode (n_tokens <= 8)
    //   * top_k values without a map0 instantiation
    //   * K < 32 (mm tile requires NK=32)
    // ADR-013 P16 — Q4_K mm_id ported; eligible for the prefill route.
    // ADR-022 Phase 2 — Q5_K mm_id ported.
    if plan_expert_auto_route(params.n_tokens, params.top_k, params.k, force_mv, routing)
        == ExpertAutoPlan::Mm
    {
        // ADR-022 AC-4: env-gated trace so operators can confirm mm_id
        // engages on prefill. `HF2Q_LOG_MM_ID_ROUTE=1` enables the line.
        if std::env::var("HF2Q_LOG_MM_ID_ROUTE").is_ok() {
            eprintln!(
                "[mlx-native adr-022 AC-4] dispatch_id_mm engaged: type={:?} \
                 n_tokens={} top_k={} k={} n={} n_experts={}",
                params.ggml_type,
                params.n_tokens,
                params.top_k,
                params.k,
                params.n,
                params.n_experts,
            );
        }
        return dispatch_id_mm(
            encoder, registry, device, input, weight, ids, output, params, routing,
        );
    }

    dispatch_id_mv(
        encoder, registry, device, input, weight, ids, output, params, routing,
    )
}

/// `output_cols` is the row width of the `[n_tokens * top_k, output_cols]`
/// output (`params.n`, or `params.n / 2` for the GEGLU epilogue).
fn validate_id_call(
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlQuantizedMatmulIdParams,
    output_cols: u32,
) -> Result<()> {
    validate_native_expert_dtypes("quantized_matmul_id_ggml", input, weight, ids, output)?;
    let qk = params.ggml_type.block_values();

    // --- Validate dimensions ---
    if params.n_tokens == 0 || params.k == 0 || params.n == 0 {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_ggml: n_tokens, K, and N must all be > 0".into(),
        ));
    }
    if params.top_k == 0 || params.top_k > params.n_experts {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_ggml: top_k must be in 1..=n_experts".into(),
        ));
    }
    if params.n_experts == 0 {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_ggml: n_experts must be > 0".into(),
        ));
    }
    if params.k % qk != 0 {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml: K ({}) must be divisible by block QK ({})",
            params.k, qk
        )));
    }

    // --- Validate buffer sizes ---
    let expected_input_bytes = checked_byte_extent(
        "input",
        &[
            params.n_tokens as usize,
            params.k as usize,
            DType::F32.size_of(),
        ],
    )?;
    if input.data_byte_len() < expected_input_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml: input buffer too small: expected {} bytes for [{} x {}] f32, got {}",
            expected_input_bytes, params.n_tokens, params.k, input.data_byte_len()
        )));
    }

    let total_weight_bytes = required_expert_weight_bytes(
        params.ggml_type,
        params.n_experts,
        params.n,
        params.k,
        params.expert_stride,
    )?;
    if weight.data_byte_len() < total_weight_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml: weight buffer too small: expected {} bytes for {} experts, got {}",
            total_weight_bytes, params.n_experts, weight.data_byte_len()
        )));
    }

    let total_rows = (params.n_tokens as usize)
        .checked_mul(params.top_k as usize)
        .ok_or_else(|| MlxError::InvalidArgument("expert row count overflow".into()))?;
    if total_rows > i32::MAX as usize {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_ggml: n_tokens * top_k exceeds signed Metal row indexing".into(),
        ));
    }
    let expected_ids_bytes = checked_byte_extent("ids", &[total_rows, DType::U32.size_of()])?;
    if ids.data_byte_len() < expected_ids_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml: ids buffer too small: expected {} bytes for [{} * {}] u32, got {}",
            expected_ids_bytes, params.n_tokens, params.top_k, ids.data_byte_len()
        )));
    }

    let expected_output_bytes = checked_byte_extent(
        "output",
        &[total_rows, output_cols as usize, DType::F32.size_of()],
    )?;
    if output.data_byte_len() < expected_output_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml: output buffer too small: expected {} bytes for [{} x {}] f32, got {}",
            expected_output_bytes, total_rows, output_cols, output.data_byte_len()
        )));
    }
    Ok(())
}

/// Same contract as `quantized_matmul_id_ggml`, but takes caller-owned
/// `IdMmScratch` so batched-prefill dispatches avoid the per-call
/// `MTLDevice.newBufferWithLength:` allocations the auto entry point
/// incurs (ADR-011 Phase 3 Wave P3b — "scratch pooling").
///
/// When the dispatch routes to the mv_id path (decode / unsupported top_k /
/// K < 32), the scratch is not touched — it is only used on the mm_id
/// path.  Callers may over-size the scratch once per prefill and share
/// it across every mm_id call in the forward pass.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_pooled(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
) -> Result<()> {
    let routing = ggml_routing_policy_for_registry(registry);
    quantized_matmul_id_ggml_pooled_with_policy(
        encoder, registry, device, input, weight, ids, output, scratch, params, &routing,
    )
}

/// Caller-scratch expert dispatch under an explicit routing policy.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_pooled_with_policy(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    quantized_matmul_id_ggml_pooled_impl(
        encoder,
        registry,
        device,
        input,
        weight,
        ids,
        output,
        scratch,
        params,
        IdMmInputLayout::SharedPerToken,
        false,
        routing,
    )
}

/// Encode two matrix-matrix expert projections that share one routing
/// schedule.
///
/// This is the family-neutral MoE gate/up primitive: both projections use the
/// same input, expert ids, quantization, and matrix shape, but may use distinct
/// weights and outputs. The first projection builds the per-expert routed-row
/// schedule; the second reuses it in the same command encoder. Their ordinary
/// `mm_id` kernels may therefore overlap after the schedule barrier without
/// duplicating `map0` or exposing schedule identity to the caller.
///
/// The pair API is intentionally matrix-matrix only. Decode-sized calls retain
/// the existing matrix-vector route and are measured separately. Output ranges
/// must not overlap because the two projection kernels may execute
/// concurrently.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_pooled_pair(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    first_weight: &MlxBuffer,
    second_weight: &MlxBuffer,
    ids: &MlxBuffer,
    first_output: &MlxBuffer,
    second_output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
) -> Result<()> {
    let routing = ggml_routing_policy_for_registry(registry);
    quantized_matmul_id_ggml_pooled_pair_with_policy(
        encoder,
        registry,
        device,
        input,
        first_weight,
        second_weight,
        ids,
        first_output,
        second_output,
        scratch,
        params,
        &routing,
    )
}

/// Paired expert-MM dispatch under an explicit routing policy.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_pooled_pair_with_policy(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    first_weight: &MlxBuffer,
    second_weight: &MlxBuffer,
    ids: &MlxBuffer,
    first_output: &MlxBuffer,
    second_output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    validate_native_expert_dtypes(
        "quantized_matmul_id_ggml_pooled_pair first projection",
        input,
        first_weight,
        ids,
        first_output,
    )?;
    validate_native_expert_dtypes(
        "quantized_matmul_id_ggml_pooled_pair second projection",
        input,
        second_weight,
        ids,
        second_output,
    )?;
    validate_id_mm_scratch_dtypes(
        "quantized_matmul_id_ggml_pooled_pair",
        &scratch.htpe,
        &scratch.hids,
    )?;
    if plan_expert_auto_route(params.n_tokens, params.top_k, params.k, false, routing)
        != ExpertAutoPlan::Mm
    {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml_pooled_pair requires the mm_id route: n_tokens={} threshold={} top_k={} k={}",
            params.n_tokens,
            routing.expert_mm_threshold,
            params.top_k,
            params.k,
        )));
    }
    if params.n == 0
        || params.n_experts == 0
        || params.top_k == 0
        || params.top_k > params.n_experts
    {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_ggml_pooled_pair requires N > 0 and top_k in 1..=n_experts".into(),
        ));
    }

    // Validate both projections' distinct buffers before encoding the first
    // projection, so a malformed pair cannot leave partial work in the
    // caller's command encoder.
    let qk = params.ggml_type.block_values();
    if params.k % qk != 0 {
        return Err(MlxError::InvalidArgument(format!(
            "K ({}) must be divisible by block QK ({})",
            params.k, qk,
        )));
    }
    let total_weight_bytes = required_expert_weight_bytes(
        params.ggml_type,
        params.n_experts,
        params.n,
        params.k,
        params.expert_stride,
    )?;
    for (name, weight) in [("first", first_weight), ("second", second_weight)] {
        if weight.data_byte_len() < total_weight_bytes {
            return Err(MlxError::InvalidArgument(format!(
                "quantized_matmul_id_ggml_pooled_pair {name} weight buffer too small: expected {} bytes, got {}",
                total_weight_bytes,
                weight.data_byte_len(),
            )));
        }
    }
    let expected_output_bytes = (params.n_tokens as usize)
        .checked_mul(params.top_k as usize)
        .and_then(|rows| rows.checked_mul(params.n as usize))
        .and_then(|elements| elements.checked_mul(DType::F32.size_of()))
        .ok_or_else(|| MlxError::InvalidArgument("pair output byte count overflow".into()))?;
    for (name, output) in [("first", first_output), ("second", second_output)] {
        if output.data_byte_len() < expected_output_bytes {
            return Err(MlxError::InvalidArgument(format!(
                "quantized_matmul_id_ggml_pooled_pair {name} output buffer too small: expected {} bytes, got {}",
                expected_output_bytes,
                output.data_byte_len(),
            )));
        }
    }

    let expected_input_bytes = (params.n_tokens as usize)
        .checked_mul(params.k as usize)
        .and_then(|elements| elements.checked_mul(DType::F32.size_of()))
        .ok_or_else(|| MlxError::InvalidArgument("pair input byte count overflow".into()))?;
    if input.data_byte_len() < expected_input_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml_pooled_pair input buffer too small: expected {} bytes, got {}",
            expected_input_bytes,
            input.data_byte_len(),
        )));
    }
    let expected_ids_bytes = (params.n_tokens as usize)
        .checked_mul(params.top_k as usize)
        .and_then(|elements| elements.checked_mul(DType::U32.size_of()))
        .ok_or_else(|| MlxError::InvalidArgument("pair ids byte count overflow".into()))?;
    if ids.data_byte_len() < expected_ids_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml_pooled_pair ids buffer too small: expected {} bytes, got {}",
            expected_ids_bytes,
            ids.data_byte_len(),
        )));
    }
    scratch.check_capacity(params.n_experts, params.n_tokens, params.top_k)?;
    let (expected_htpe_bytes, expected_hids_bytes) = id_mm_scratch_bytes(
        routing.expert_mm_schedule,
        params.n_experts,
        params.n_tokens,
        params.top_k,
    )?;
    for (name, buffer, expected) in [
        ("htpe", &scratch.htpe, expected_htpe_bytes),
        ("hids", &scratch.hids, expected_hids_bytes),
    ] {
        if buffer.data_byte_len() < expected {
            return Err(MlxError::InvalidArgument(format!(
                "quantized_matmul_id_ggml_pooled_pair {name} scratch buffer too small: expected {} bytes, got {}",
                expected,
                buffer.data_byte_len(),
            )));
        }
    }

    let range = |buffer: &MlxBuffer, extent: usize| {
        let start = (buffer.contents_ptr() as usize).saturating_add(buffer.byte_offset() as usize);
        (start, start.saturating_add(extent))
    };
    let overlaps =
        |left: (usize, usize), right: (usize, usize)| left.0 < right.1 && right.0 < left.1;
    let first_output_range = range(first_output, expected_output_bytes);
    let second_output_range = range(second_output, expected_output_bytes);
    if overlaps(first_output_range, second_output_range) {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_ggml_pooled_pair output ranges must not overlap".into(),
        ));
    }
    let scratch_ranges = [
        ("htpe", range(&scratch.htpe, expected_htpe_bytes)),
        ("hids", range(&scratch.hids, expected_hids_bytes)),
    ];
    if overlaps(scratch_ranges[0].1, scratch_ranges[1].1) {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_ggml_pooled_pair scratch ranges must not overlap".into(),
        ));
    }
    let immutable_read_ranges = [
        ("input", range(input, expected_input_bytes)),
        ("first weight", range(first_weight, total_weight_bytes)),
        ("second weight", range(second_weight, total_weight_bytes)),
        ("ids", range(ids, expected_ids_bytes)),
    ];
    for (scratch_name, scratch_range) in scratch_ranges {
        if let Some((read_name, _)) = immutable_read_ranges
            .iter()
            .find(|(_, read_range)| overlaps(scratch_range, *read_range))
        {
            return Err(MlxError::InvalidArgument(format!(
                "quantized_matmul_id_ggml_pooled_pair {scratch_name} scratch range must not overlap {read_name}",
            )));
        }
    }
    let read_ranges = [
        immutable_read_ranges[0],
        immutable_read_ranges[1],
        immutable_read_ranges[2],
        immutable_read_ranges[3],
        ("htpe scratch", scratch_ranges[0].1),
        ("hids scratch", scratch_ranges[1].1),
    ];
    for (output_name, output_range) in [
        ("first", first_output_range),
        ("second", second_output_range),
    ] {
        if let Some((read_name, _)) = read_ranges
            .iter()
            .find(|(_, read_range)| overlaps(output_range, *read_range))
        {
            return Err(MlxError::InvalidArgument(format!(
                "quantized_matmul_id_ggml_pooled_pair {output_name} output range must not overlap {read_name}",
            )));
        }
    }

    dispatch_id_mm_pooled_with_layout(
        encoder,
        registry,
        device,
        input,
        first_weight,
        ids,
        first_output,
        scratch,
        params,
        IdMmInputLayout::SharedPerToken,
        false,
        routing,
    )?;
    dispatch_id_mm_pooled_with_layout(
        encoder,
        registry,
        device,
        input,
        second_weight,
        ids,
        second_output,
        scratch,
        params,
        IdMmInputLayout::SharedPerToken,
        true,
        routing,
    )
}

/// Pooled `mm_id` dispatch for input laid out as
/// `[n_tokens, top_k, K]` instead of one shared row per token.
///
/// This serves MoE down projections whose activation already has one
/// distinct row per routed expert slot. The Metal `mm_id` kernels natively
/// support the slot and token strides, so preserving the shape avoids
/// flattening `top_k` into the token dimension and over-dispatching the
/// expert grid by that factor.
///
/// The slotted contract is matrix-matrix only. Calls at or below the `mm_id`
/// routing threshold are rejected instead of entering the flat `mv_id` path
/// with incompatible input addressing.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_pooled_slotted(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
) -> Result<()> {
    let routing = ggml_routing_policy_for_registry(registry);
    quantized_matmul_id_ggml_pooled_slotted_with_policy(
        encoder, registry, device, input, weight, ids, output, scratch, params, &routing,
    )
}

/// Slotted caller-scratch expert MM under an explicit routing policy.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_pooled_slotted_with_policy(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    quantized_matmul_id_ggml_pooled_impl(
        encoder,
        registry,
        device,
        input,
        weight,
        ids,
        output,
        scratch,
        params,
        IdMmInputLayout::Slotted,
        false,
        routing,
    )
}

/// Slotted pooled `mm_id` that consumes the routing schedule already in
/// `scratch` instead of encoding `map0` (ADR-063 P2b).
///
/// Intended for an MoE down projection whose gate/up call built the
/// schedule from the same expert ids: the down input is `[n_tokens, top_k,
/// K]`, the ids are the gate/up call's `[n_tokens, top_k]` table, and the
/// schedule depends only on the ids, so the result is byte-identical to
/// [`quantized_matmul_id_ggml_pooled_slotted`].
///
/// Every `map0` encode into `scratch` records a schedule key: the ids
/// buffer (object, byte offset and extent), `n_tokens`, `top_k`,
/// `n_experts`, the schedule kind, the [`CommandEncoder`] object (a
/// `commit_wait_and_rotate` keeps it) and a per-scratch encode serial. A
/// mutable buffer access through [`IdMmScratch::htpe_mut`],
/// [`IdMmScratch::hids_mut`] or [`IdMmScratch::buffers_mut`] clears the key.
/// This call returns `InvalidArgument` before encoding anything unless the
/// recorded key equals the key of this call (same ids, shape, schedule and
/// encoder, no `map0` or mutable access since). The `map0` barrier already
/// in the encoder orders schedule construction before this dispatch.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_pooled_slotted_reusing_schedule(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
) -> Result<()> {
    let routing = ggml_routing_policy_for_registry(registry);
    quantized_matmul_id_ggml_pooled_slotted_reusing_schedule_with_policy(
        encoder, registry, device, input, weight, ids, output, scratch, params, &routing,
    )
}

/// Schedule-reusing slotted expert MM under an explicit routing policy.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_ggml_pooled_slotted_reusing_schedule_with_policy(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    quantized_matmul_id_ggml_pooled_impl(
        encoder,
        registry,
        device,
        input,
        weight,
        ids,
        output,
        scratch,
        params,
        IdMmInputLayout::Slotted,
        true,
        routing,
    )
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum IdMmInputLayout {
    SharedPerToken,
    Slotted,
}

#[allow(clippy::too_many_arguments)]
fn quantized_matmul_id_ggml_pooled_impl(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
    input_layout: IdMmInputLayout,
    reuse_schedule: bool,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    validate_native_expert_dtypes(
        "quantized_matmul_id_ggml_pooled",
        input,
        weight,
        ids,
        output,
    )?;
    validate_id_mm_scratch_dtypes(
        "quantized_matmul_id_ggml_pooled",
        &scratch.htpe,
        &scratch.hids,
    )?;
    // Mirror the validation + routing logic from `quantized_matmul_id_ggml`
    // so the pooled path has identical correctness invariants.  (We keep
    // the two entry points separate rather than extracting a shared inner
    // because the scratch is only relevant on the mm_id branch — lifting
    // scratch into the mv branch would add unused parameters.)
    let qk = params.ggml_type.block_values();

    if params.n_tokens == 0 || params.k == 0 || params.n == 0 {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_ggml_pooled: n_tokens, K, and N must all be > 0".into(),
        ));
    }
    if params.top_k == 0 || params.top_k > params.n_experts {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_ggml_pooled: top_k must be in 1..=n_experts".into(),
        ));
    }
    if params.n_experts == 0 {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_ggml_pooled: n_experts must be > 0".into(),
        ));
    }
    if params.k % qk != 0 {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml_pooled: K ({}) must be divisible by block QK ({})",
            params.k, qk
        )));
    }

    let input_rows = match input_layout {
        IdMmInputLayout::SharedPerToken => params.n_tokens as usize,
        IdMmInputLayout::Slotted => (params.n_tokens as usize)
            .checked_mul(params.top_k as usize)
            .ok_or_else(|| MlxError::InvalidArgument("slotted input row count overflow".into()))?,
    };
    let expected_input_bytes = input_rows
        .checked_mul(params.k as usize)
        .and_then(|elements| elements.checked_mul(DType::F32.size_of()))
        .ok_or_else(|| MlxError::InvalidArgument("input byte count overflow".into()))?;
    if input.data_byte_len() < expected_input_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml_pooled: input buffer too small: expected {} bytes for [{} x {}] f32 {:?} input, got {}",
            expected_input_bytes, input_rows, params.k, input_layout, input.data_byte_len()
        )));
    }

    let total_weight_bytes = required_expert_weight_bytes(
        params.ggml_type,
        params.n_experts,
        params.n,
        params.k,
        params.expert_stride,
    )?;
    if weight.data_byte_len() < total_weight_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml_pooled: weight buffer too small: expected {} bytes for {} experts, got {}",
            total_weight_bytes, params.n_experts, weight.data_byte_len()
        )));
    }

    let total_rows = (params.n_tokens as usize)
        .checked_mul(params.top_k as usize)
        .ok_or_else(|| MlxError::InvalidArgument("expert row count overflow".into()))?;
    if total_rows > i32::MAX as usize {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_ggml_pooled: n_tokens * top_k exceeds signed Metal row indexing"
                .into(),
        ));
    }
    let expected_ids_bytes = checked_byte_extent("ids", &[total_rows, DType::U32.size_of()])?;
    if ids.data_byte_len() < expected_ids_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml_pooled: ids buffer too small: expected {} bytes for [{} * {}] u32, got {}",
            expected_ids_bytes, params.n_tokens, params.top_k, ids.data_byte_len()
        )));
    }

    let expected_output_bytes = checked_byte_extent(
        "output",
        &[total_rows, params.n as usize, DType::F32.size_of()],
    )?;
    if output.data_byte_len() < expected_output_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml_pooled: output buffer too small: expected {} bytes for [{} x {}] f32, got {}",
            expected_output_bytes, total_rows, params.n, output.data_byte_len()
        )));
    }

    // Accept the exact map0 template set {1, 6, 8}; top-k 6 serves DeepSeek-V4.
    // ADR-013 P16 — Q4_K mm_id ported; eligible for the prefill route.
    // ADR-022 Phase 2 — Q5_K mm_id ported.  Previous bypass here
    // is retired (kernels live in id_mm.metal + id_mm_tensor.metal).
    // ADR-022 AC-4: env-gated trace (HF2Q_LOG_MM_ID_ROUTE=1) confirms mm_id
    // engagement on the qwen35 prefill path which goes through this pooled
    // entry, not the auto entry above.
    if plan_expert_auto_route(params.n_tokens, params.top_k, params.k, false, routing)
        == ExpertAutoPlan::Mm
    {
        if std::env::var("HF2Q_LOG_MM_ID_ROUTE").is_ok() {
            eprintln!(
                "[mlx-native adr-022 AC-4 pooled] dispatch_id_mm_pooled engaged: \
                 type={:?} n_tokens={} top_k={} k={} n={} n_experts={}",
                params.ggml_type,
                params.n_tokens,
                params.top_k,
                params.k,
                params.n,
                params.n_experts,
            );
        }
        if reuse_schedule {
            scratch.validate_schedule_reuse(
                encoder,
                ids,
                params.n_tokens,
                params.top_k,
                params.n_experts,
                routing.expert_mm_schedule,
            )?;
        }
        return dispatch_id_mm_pooled_with_layout(
            encoder,
            registry,
            device,
            input,
            weight,
            ids,
            output,
            scratch,
            params,
            input_layout,
            reuse_schedule,
            routing,
        );
    }

    if input_layout == IdMmInputLayout::Slotted {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_ggml_pooled_slotted requires the mm_id route: n_tokens={} threshold={} top_k={} k={}",
            params.n_tokens, routing.expert_mm_threshold, params.top_k, params.k,
        )));
    }

    dispatch_id_mv(
        encoder, registry, device, input, weight, ids, output, params, routing,
    )
}

/// The n_tokens threshold at which `quantized_matmul_id_ggml` switches
/// from the mv_id kernel to the mm_id kernel.  Matches the reference
/// `ne11_mm_min = 8`.
// ADR-013 P17 — bumped 8 → 32 to match the reference `ne21_mm_id_min = 32`.
// Below 32 tokens, mv_id is faster than mm_id
// (the mm tile-reuse setup overhead doesn't amortize at small n_tokens).
pub const MM_ID_ROUTING_THRESHOLD: u32 = 32;

/// Threadgroup memory for `kernel_mul_mv_id_iq2_xxs_f32`: the 256-entry
/// `iq2xxs_grid` (256 × 8 B) followed by the 128 `ksigns_iq2xs` bytes.
const IQ2_XXS_MV_ID_SHARED_BYTES: u64 = 256 * 8 + 128;

/// ADR-013 P19 H11 (2026-05-01) — runtime override for the mm_id routing
/// threshold via `HF2Q_MM_ID_ROUTING_THRESHOLD` env. Setting it to a very
/// large value (e.g. `99999`) forces every dispatch onto the `mv_id` route,
/// which falsifies / confirms the "mm_id setup overhead is dominant at
/// pp ≤ 256" hypothesis without recompiling. Read once into a `OnceLock`
/// at first call and reused; absent or unparseable env falls back to the
/// compile-time const above.
fn mm_id_routing_threshold() -> u32 {
    static CACHED: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
    *CACHED.get_or_init(|| {
        std::env::var("HF2Q_MM_ID_ROUTING_THRESHOLD")
            .ok()
            .and_then(|s| s.parse::<u32>().ok())
            .map(|v| {
                if std::env::var("MLX_LOG_TENSOR_PROBE").is_ok() {
                    eprintln!(
                        "[mlx-native] mm_id_routing_threshold: OVERRIDE via HF2Q_MM_ID_ROUTING_THRESHOLD={v} (default {})",
                        MM_ID_ROUTING_THRESHOLD
                    );
                }
                v
            })
            .unwrap_or(MM_ID_ROUTING_THRESHOLD)
    })
}

fn mv_id_gpu_params(params: &GgmlQuantizedMatmulIdParams) -> GgmlMatvecIdGpuParams {
    GgmlMatvecIdGpuParams {
        ne00: params.k as i64,
        ne01: params.n as i64,
        ne02: params.n_experts as i64,
        ne10: params.k as i64,
        ne12: 1,
        ne0: params.n as i64,
        ne1: (params.n_tokens as i64) * (params.top_k as i64),
        r2: 1,
        r3: 1,
        top_k: params.top_k,
        n_tokens: params.n_tokens,
        expert_stride: params.expert_stride as i64,
    }
}

/// Matrix-vector `_id` dispatch (decode path, unchanged from pre-Phase-3).
#[allow(clippy::too_many_arguments)]
fn dispatch_id_mv(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlQuantizedMatmulIdParams,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    let total_rows = (params.n_tokens as usize) * (params.top_k as usize);

    // ADR-028 —nr0=2 variant for q6_K _id mat-vec.  Mirrors
    // the non-_id work.
    //
    // ADR-028 default-flipped to ON (operator REFRAME #2).
    // Opt out with `HF2Q_Q6K_ID_MV_NR2=0` / `=false` / `=off`.
    let use_q6k_id_nr2 = matches!(params.ggml_type, GgmlType::Q6_K) && routing.expert_q6k_mv_nr2;
    // ADR-029 — nr0=2 nsg=4 variant for q8_0 _id mat-vec.
    // Matches peer's N_R0_Q8_0=2 + N_SG_Q8_0=4.
    // gemma4 APEX-Q5_K_M MoE down_exps is Q8_0 → 30 dispatches/decode-tok
    // on this path.  Opt-in via `HF2Q_Q8_0_ID_MV_NR2=1`; default-off
    // until coherence + bench validation. Cached via AtomicI8.
    let use_q8_0_id_nr2 = matches!(params.ggml_type, GgmlType::Q8_0) && routing.expert_q8_0_mv_nr2;
    let kernel_name = if use_q6k_id_nr2 {
        "kernel_mul_mv_id_q6_K_f32_nr2"
    } else if use_q8_0_id_nr2 {
        "kernel_mul_mv_id_q8_0_f32_nr2"
    } else {
        params.ggml_type.id_kernel_name()
    };
    let pipeline = registry.get_pipeline(kernel_name, device.metal_device())?;

    let gpu_params = mv_id_gpu_params(params);

    let (nth0, nth1, align) = match params.ggml_type {
        // Q4_0/Q8_0: historical (8, 8) layout = 64 threads = 2 simdgroups
        // of 32.  Tested 2026-04-26 against the reference (32, 2) layout —
        // (32, 2) gave a 1.8% short-bench improvement on dwq46 but a
        // 2.0% REGRESSION on the 5-cold-run 256-token decode bench
        // (median 108.0 vs 110.2 t/s).  Likely cache/scheduling
        // interaction with the multi-token KV-cache memory access
        // pattern that the (8, 8) layout happens to align with on M5
        // Max.  6th confirmed M5 Max static-evidence kernel hypothesis
        // falsified — Metal compiler/scheduler optimizes both layouts
        // similarly, with workload-specific edges that don't match
        // the reference tuning.
        // ADR-022: Q5_1 and IQ4_NL are 32-element legacy formats; share
        // the (8, 8) layout with Q4_0 / Q8_0. Confirmed against the reference
        // dispatch_id_mv launch geometry for `kernel_mul_mv_id_q5_1_f32`
        // and `kernel_mul_mv_id_iq4_nl_f32` (both NWG=2, NSIMDGROUP=2,
        // ngroups along K = nb/4 → 8 thread blocks of 8 rows each).
        GgmlType::Q4_0
        | GgmlType::Q5_0
        | GgmlType::Q8_0
        | GgmlType::Q5_1
        | GgmlType::IQ4_NL
        // ADR-033 §Pi Task #16 — IQ4_XS mv_id mirrors IQ4_NL's
        // (N_SIMDGROUP=2, N_DST=4, NWG=2) launch geometry.
        | GgmlType::IQ4_XS => (8u64, 8u64, 8usize),
        GgmlType::Q2_K => (2u64, 32u64, 8usize),
        GgmlType::Q3_K => (2u64, 32u64, 4usize),
        // Q4_K, Q5_K, and Q6_K all use the 2-row-per-threadgroup (2, 32)
        // geometry.  ADR-013 P7 — Q4_K added; mirrors Q5_K (NSG=2,
        // 1 row per simdgroup; same kmask scale-decode).
        GgmlType::Q4_K | GgmlType::Q5_K | GgmlType::Q6_K => (2u64, 32u64, 2usize),
        // ADR-066 N-IQ2XXS — two simdgroups of 32 lanes, four rows each.
        GgmlType::IQ2_XXS => (32u64, 2u64, 8usize),
        GgmlType::F32
        | GgmlType::F16
        | GgmlType::BF16
        | GgmlType::I16
        | GgmlType::I32 => {
            return Err(MlxError::InvalidArgument(format!(
                "quantized_matmul_id_ggml does not support {:?}",
                params.ggml_type
            )));
        }
    };
    // ADR-028 —nr0=2 doubles rows-per-TG to 4. Same 2 SGs × 32
    // threads, but each SG handles 2 rows so align=4.
    let align = if use_q6k_id_nr2 { 4usize } else { align };
    // ADR-029 — Q8_0 _id NR2 NSG=4: threads_per_tg=(32, 4), align=2
    // (each TG covers NR0=2 rows; all 4 SGs collaborate on K-dim).
    // Override geometry AFTER the Q4_0/Q8_0/Q5_1/IQ4_NL match arm above.
    let (nth0, nth1, align) = if use_q8_0_id_nr2 {
        (32u64, 4u64, 2usize)
    } else {
        (nth0, nth1, align)
    };

    let n = params.n as usize;
    let m = total_rows;

    // Dispatch routing dim in Y, NOT Z (despite the reference mul_mv_id
    // routing in z).  Tested 2026-04-26: switching
    // to z-routing on M5 Max regressed dwq46 256-token decode from 112 t/s
    // to 90.9 t/s (-19%).  Apple GPU's threadgroup scheduler distributes
    // this dispatch shape better via y than z — 7th confirmed static-
    // evidence kernel hypothesis falsified per
    // `project_metal_compiler_auto_optimizes_static_levers.md`.
    let threadgroups = metal::MTLSize::new(div_ceil(n, align) as u64, m as u64, 1);
    let threads_per_tg = metal::MTLSize::new(nth0, nth1, 1);

    if matches!(params.ggml_type, GgmlType::IQ2_XXS) {
        // The IQ2_XXS kernel stages the 256-entry grid (2,048 B) and the
        // 128 sign bytes in threadgroup memory once per threadgroup.
        encoder.encode_threadgroups_with_args_and_shared(
            pipeline,
            &[
                (0, KernelArg::Buffer(weight)),
                (1, KernelArg::Buffer(input)),
                (2, KernelArg::Buffer(output)),
                (3, KernelArg::Buffer(ids)),
                (4, KernelArg::Bytes(as_bytes(&gpu_params))),
            ],
            &[(0, IQ2_XXS_MV_ID_SHARED_BYTES)],
            threadgroups,
            threads_per_tg,
        );
    } else if use_q8_0_id_nr2 {
        // ADR-029: cross-SG reduction needs threadgroup memory:
        // NR0 * NW * sizeof(float) = 2 * 32 * 4 = 256 bytes.
        let smem_bytes: u64 = 2 * 32 * std::mem::size_of::<f32>() as u64;
        encoder.encode_threadgroups_with_args_and_shared(
            pipeline,
            &[
                (0, KernelArg::Buffer(weight)),
                (1, KernelArg::Buffer(input)),
                (2, KernelArg::Buffer(output)),
                (3, KernelArg::Buffer(ids)),
                (4, KernelArg::Bytes(as_bytes(&gpu_params))),
            ],
            &[(0, smem_bytes)],
            threadgroups,
            threads_per_tg,
        );
    } else {
        encoder.encode_threadgroups_with_args(
            pipeline,
            &[
                (0, KernelArg::Buffer(weight)),
                (1, KernelArg::Buffer(input)),
                (2, KernelArg::Buffer(output)),
                (3, KernelArg::Buffer(ids)),
                (4, KernelArg::Bytes(as_bytes(&gpu_params))),
            ],
            threadgroups,
            threads_per_tg,
        );
    }

    Ok(())
}

/// ADR-029 — pre-bake the per-weight Q6_K `_id` NR2
/// `m=1` decode dispatch into a `DispatchRecord`.
///
/// The MoE gate_up dispatch in gemma4 APEX-Q5_K_M is Q6_K, hits ~30
/// calls/decode-tok (1 dispatch per layer, n_tokens=1, total_rows=top_k
/// folded into `threadgroups.y` per `dispatch_id_mv`'s geometry).  Same
/// rationale as `quantized_matmul_ggml::build_q6k_nr2_m1_record` — bake
/// the load-time-immutable parts of the call once at first-dispatch.
///
/// Pre-bakes:
///   - Pipeline reference (skips `KernelRegistry` HashMap lookup per call)
///   - MTLSize threadgroups + threads_per_tg (skips MTLSize::new + match)
///   - `GgmlMatvecIdGpuParams` bytes (skips struct construction + bytemuck)
///   - Binding slot order: weight=0, input=1, output=2, ids=3, params=4
///
/// Returns `None` when `HF2Q_Q6K_ID_MV_NR2` is set off — the
/// non-NR2 kernel uses a different geometry (2 SGs × 1 row/SG = align=2)
/// and a record baked for the NR2 layout would be wrong.  Callers MUST
/// fall through to the unbaked `dispatch_id_mv` path in that case.
///
/// Bake-time validation: pipeline lookup must succeed.  Threadgroup
/// geometry is hard-coded to the Q6_K_ID NR2 contract (NSG=2, nr0=2,
/// align=4 rows/TG, threads=(2, 32, 1)) — mirrors the kernel
/// `kernel_mul_mv_id_q6_K_f32_nr2`'s `constexpr int NSG = 2; constexpr
/// int nr0 = 2;` at `quantized_matmul_id_ggml.metal:1025-1026`.
///
/// `top_k`, `n_experts`, and `expert_stride` are weight-specific and folded into the
/// baked `GgmlMatvecIdGpuParams.ne1 = top_k` (since `n_tokens=1` at
/// decode), the on-device expert-ID bound, and `expert_stride`. Callers
/// must pass the same values they would use in
/// `GgmlQuantizedMatmulIdParams`.
pub fn build_q6k_id_nr2_m1_record(
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    n: u32,
    k: u32,
    top_k: u32,
    n_experts: u32,
    expert_stride: u64,
) -> Result<Option<DispatchRecord>> {
    let routing = ggml_routing_policy_for_registry(registry);
    build_q6k_id_nr2_m1_record_with_policy(
        registry,
        device,
        n,
        k,
        top_k,
        n_experts,
        expert_stride,
        &routing,
    )
}

/// Explicit-policy form of [`build_q6k_id_nr2_m1_record`].
pub fn build_q6k_id_nr2_m1_record_with_policy(
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    n: u32,
    k: u32,
    top_k: u32,
    n_experts: u32,
    expert_stride: u64,
    routing: &GgmlRoutingPolicy,
) -> Result<Option<DispatchRecord>> {
    if n_experts == 0 || top_k == 0 || top_k > n_experts {
        return Err(MlxError::InvalidArgument(
            "build_q6k_id_nr2_m1_record requires top_k in 1..=n_experts".into(),
        ));
    }
    if top_k > i32::MAX as u32 {
        return Err(MlxError::InvalidArgument(
            "build_q6k_id_nr2_m1_record top_k exceeds signed Metal row indexing".into(),
        ));
    }
    if expert_stride > i64::MAX as u64 {
        return Err(MlxError::InvalidArgument(
            "expert stride exceeds the signed Metal kernel ABI".into(),
        ));
    }
    // Only bakeable when the NR2 variant is the selected one — same
    // gate as `dispatch_id_mv`'s `use_q6k_id_nr2` branch.
    if !routing.expert_q6k_mv_nr2 {
        return Ok(None);
    }

    // Pipeline lookup — `kernel_mul_mv_id_q6_K_f32_nr2` takes no
    // function constants (unlike the non-_id Q6_K NR2 which carries
    // 700/701/702 for ne12/r2/r3 promotion).
    let pipeline = registry
        .get_pipeline("kernel_mul_mv_id_q6_K_f32_nr2", device)?
        .clone();

    // GgmlMatvecIdGpuParams for n_tokens=1.  total_rows = 1 * top_k.
    let gpu_params = GgmlMatvecIdGpuParams {
        ne00: k as i64,
        ne01: n as i64,
        ne02: n_experts as i64,
        ne10: k as i64,
        ne12: 1,
        ne0: n as i64,
        ne1: top_k as i64,
        r2: 1,
        r3: 1,
        top_k,
        n_tokens: 1,
        expert_stride: expert_stride as i64,
    };
    let params_bytes = as_bytes(&gpu_params).to_vec();

    // Q6_K_ID NR2: align=4 rows per TG, threads = (nth0=2, nth1=32, 1)
    // (matches `dispatch_id_mv`'s Q6_K NR2 branch override).
    const ALIGN: u32 = 4;
    let threadgroups =
        metal::MTLSize::new(div_ceil(n as usize, ALIGN as usize) as u64, top_k as u64, 1);
    let threads_per_tg = metal::MTLSize::new(2, 32, 1);

    Ok(Some(DispatchRecord {
        pipeline,
        threadgroups,
        threads_per_tg,
        threadgroup_mem: Vec::new(), // Q6_K_ID NR2 doesn't use shmem
        params_bytes,
        params_slot: 4,
        buffer_slots: vec![0, 1, 2, 3], // weight, input, output, ids
        op_kind: CapturedOpKind::Other,
        kernel_name: "kernel_mul_mv_id_q6_K_f32_nr2".to_string(),
    }))
}

/// ADR-029 — pre-bake the per-weight Q8_0 `_id` (regular,
/// non-NR2) decode dispatch into a `DispatchRecord`.
///
/// The MoE down dispatch in gemma4 APEX-Q5_K_M is Q8_0 → ~30 calls/decode-tok
/// (1 down dispatch per layer × 30 layers).  Unlike the gate_up path, the
/// down call site passes `n_tokens = real_top_k, top_k = 1` to
/// `quantized_matmul_id_ggml` — so the kernel sees `total_rows = real_top_k`
/// folded into `threadgroups.y`, and `params.ne1 = real_top_k`.  At decode
/// time `real_top_k` is fixed by model config, so the bake is valid.
///
/// Pre-bakes:
///   - Pipeline reference for `kernel_mul_mv_id_q8_0_f32` (regular variant)
///   - threadgroups = `(div_ceil(n, 8), real_top_k, 1)`
///   - threads_per_tg = `(8, 8, 1)` (N_DST=4 × N_SIMDGROUP=2 contract;
///     2 SGs × 32 threads = 64 threads/TG; 8 rows/TG)
///   - `GgmlMatvecIdGpuParams` bytes (n_tokens=real_top_k, top_k=1, ne1=real_top_k)
///   - Binding slot order: weight=0, input=1, output=2, ids=3, params=4
///
/// Returns `None` when `HF2Q_Q8_0_ID_MV_NR2` is on — the opt-in NR2 kernel
/// uses different geometry (threads=(32, 4, 1), align=2, shmem=256 bytes)
/// and a record baked for the regular layout would be wrong.  Callers MUST
/// fall through to the unbaked `dispatch_id_mv` path in that case.
///
/// Bake-time validation: pipeline lookup must succeed.  Geometry is
/// hard-coded to the regular Q8_0_ID contract (N_DST=4, N_SIMDGROUP=2,
/// align=8) — mirrors the kernel `kernel_mul_mv_id_q8_0_f32` constants
/// at `quantized_matmul_id_ggml.metal:460-462` and the dispatch_id_mv
/// `(nth0=8, nth1=8, align=8)` branch.
///
/// `real_top_k`, `n_experts`, and `expert_stride` are weight-specific; callers must
/// pass the same `top_k` value the model config carries (folded into
/// `params.n_tokens` at the down call site), the same expert count used for
/// the on-device ID bound, and the same `expert_stride` they would have used
/// in `GgmlQuantizedMatmulIdParams`.
pub fn build_q8_0_id_decode_record(
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    n: u32,
    k: u32,
    real_top_k: u32,
    n_experts: u32,
    expert_stride: u64,
) -> Result<Option<DispatchRecord>> {
    let routing = ggml_routing_policy_for_registry(registry);
    build_q8_0_id_decode_record_with_policy(
        registry,
        device,
        n,
        k,
        real_top_k,
        n_experts,
        expert_stride,
        &routing,
    )
}

/// Explicit-policy form of [`build_q8_0_id_decode_record`].
pub fn build_q8_0_id_decode_record_with_policy(
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    n: u32,
    k: u32,
    real_top_k: u32,
    n_experts: u32,
    expert_stride: u64,
    routing: &GgmlRoutingPolicy,
) -> Result<Option<DispatchRecord>> {
    if n_experts == 0 || real_top_k == 0 || real_top_k > n_experts {
        return Err(MlxError::InvalidArgument(
            "build_q8_0_id_decode_record requires real_top_k in 1..=n_experts".into(),
        ));
    }
    if real_top_k > i32::MAX as u32 {
        return Err(MlxError::InvalidArgument(
            "build_q8_0_id_decode_record real_top_k exceeds signed Metal row indexing".into(),
        ));
    }
    if expert_stride > i64::MAX as u64 {
        return Err(MlxError::InvalidArgument(
            "expert stride exceeds the signed Metal kernel ABI".into(),
        ));
    }
    // Only bakeable when the regular (non-NR2) Q8_0_ID kernel is selected.
    // Mirrors the negation of `dispatch_id_mv`'s `use_q8_0_id_nr2` branch.
    if routing.expert_q8_0_mv_nr2 {
        return Ok(None);
    }

    // Pipeline lookup — regular `kernel_mul_mv_id_q8_0_f32` takes no
    // function constants.
    let pipeline = registry
        .get_pipeline("kernel_mul_mv_id_q8_0_f32", device)?
        .clone();

    // `GgmlMatvecIdGpuParams` for the down dispatch: n_tokens=real_top_k,
    // top_k=1, total_rows = n_tokens * top_k = real_top_k.
    let gpu_params = GgmlMatvecIdGpuParams {
        ne00: k as i64,
        ne01: n as i64,
        ne02: n_experts as i64,
        ne10: k as i64,
        ne12: 1,
        ne0: n as i64,
        ne1: real_top_k as i64,
        r2: 1,
        r3: 1,
        top_k: 1,
        n_tokens: real_top_k,
        expert_stride: expert_stride as i64,
    };
    let params_bytes = as_bytes(&gpu_params).to_vec();

    // Regular Q8_0_ID: align=8 rows per TG, threads=(8, 8, 1).
    const ALIGN: u32 = 8;
    let threadgroups = metal::MTLSize::new(
        div_ceil(n as usize, ALIGN as usize) as u64,
        real_top_k as u64,
        1,
    );
    let threads_per_tg = metal::MTLSize::new(8, 8, 1);

    Ok(Some(DispatchRecord {
        pipeline,
        threadgroups,
        threads_per_tg,
        threadgroup_mem: Vec::new(), // regular Q8_0_ID doesn't use shmem
        params_bytes,
        params_slot: 4,
        buffer_slots: vec![0, 1, 2, 3], // weight, input, output, ids
        op_kind: CapturedOpKind::Other,
        kernel_name: "kernel_mul_mv_id_q8_0_f32".to_string(),
    }))
}

/// Fused SwiGLU + expert-routed Q4_0 mat-vec.
///
/// Computes `output[r][n] = sum_k(dequant(W_q4_0[ids[r]][n][k]) * (silu(gate[r][k]) * up[r][k]))`
/// in a single dispatch — replaces the `silu_mul + quantized_matmul_id_ggml`
/// sequence that hf2q's MoE FFN decode path used (Phase D + Phase E in
/// `gpu_ffn.rs:build_moe_ffn_layer_gpu_q_into`).
///
/// Saves one dispatch + one memory_barrier per MoE layer × 40 layers per
/// decode token, targeting the dispatch-count component of the dwq46
/// 0.93× decode parity gap (per the `MLX_PROFILE_CB=1` per-cb breakdown
/// in ADR-012 §Optimize / Task #15).
///
/// Currently supports `GgmlType::Q4_0` only (the dominant expert-down
/// quant type in dwq46).  Q8_0 / Q6_K support is straightforward to add
/// by templating the inner dot-product over the dequant kernel; not yet
/// implemented because dwq46's expert_down is 95% Q4_0.
///
/// # Arguments
///
/// * `encoder`  -- Command encoder.
/// * `registry` -- Kernel registry.
/// * `device`   -- Metal device.
/// * `gate`     -- f32 gate input `[n_tokens*top_k, K]`.
/// * `up`       -- f32 up input `[n_tokens*top_k, K]`.
/// * `weight`   -- Q4_0 expert weight stack `[n_experts, N, packed_K]`.
/// * `ids`      -- u32 expert index buffer `[n_tokens*top_k]`.
/// * `output`   -- f32 output `[n_tokens*top_k, N]`.
/// * `params`   -- Same params struct as `quantized_matmul_id_ggml` (must have ggml_type=Q4_0).
///
/// # Errors
///
/// Returns `MlxError::InvalidArgument` if dimensions don't match expected
/// shapes or `params.ggml_type` is not `Q4_0`.
#[allow(clippy::too_many_arguments)]
pub fn quantized_matmul_id_swiglu_q4_0(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    gate: &MlxBuffer,
    up: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlQuantizedMatmulIdParams,
) -> Result<()> {
    if gate.dtype() != DType::F32
        || up.dtype() != DType::F32
        || weight.dtype() != DType::U8
        || ids.dtype() != DType::U32
        || output.dtype() != DType::F32
    {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_swiglu_q4_0 requires F32 gate/up, native U8 GGUF blocks, U32 expert IDs, and F32 output; got {:?}/{:?}/{:?}/{:?}/{:?}",
            gate.dtype(),
            up.dtype(),
            weight.dtype(),
            ids.dtype(),
            output.dtype(),
        )));
    }
    if params.ggml_type != GgmlType::Q4_0 {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_swiglu_q4_0: expected Q4_0, got {:?}",
            params.ggml_type
        )));
    }
    let qk = GgmlType::Q4_0.block_values();

    // --- Validate dimensions (mirror quantized_matmul_id_ggml) ---
    if params.n_tokens == 0 || params.k == 0 || params.n == 0 || params.n_experts == 0 {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_swiglu_q4_0: n_tokens, K, N, and n_experts must all be > 0".into(),
        ));
    }
    if params.top_k == 0 || params.top_k > params.n_experts {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_swiglu_q4_0: top_k must be in 1..=n_experts".into(),
        ));
    }
    if params.k % qk != 0 {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_swiglu_q4_0: K ({}) must be divisible by block QK ({})",
            params.k, qk
        )));
    }

    let expected_weight_bytes = required_expert_weight_bytes(
        params.ggml_type,
        params.n_experts,
        params.n,
        params.k,
        params.expert_stride,
    )?;
    if weight.data_byte_len() < expected_weight_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_swiglu_q4_0: weight buffer too small: expected {} bytes, got {}",
            expected_weight_bytes,
            weight.data_byte_len(),
        )));
    }

    let total_rows = (params.n_tokens as usize)
        .checked_mul(params.top_k as usize)
        .ok_or_else(|| MlxError::InvalidArgument("expert row count overflow".into()))?;
    if total_rows > i32::MAX as usize {
        return Err(MlxError::InvalidArgument(
            "quantized_matmul_id_swiglu_q4_0: n_tokens * top_k exceeds signed Metal row indexing"
                .into(),
        ));
    }
    let expected_in_bytes = checked_byte_extent(
        "swiglu input",
        &[total_rows, params.k as usize, DType::F32.size_of()],
    )?;
    if gate.data_byte_len() < expected_in_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_swiglu_q4_0: gate buffer too small: expected {} bytes, got {}",
            expected_in_bytes,
            gate.data_byte_len()
        )));
    }
    if up.data_byte_len() < expected_in_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_swiglu_q4_0: up buffer too small: expected {} bytes, got {}",
            expected_in_bytes,
            up.data_byte_len()
        )));
    }
    let expected_out_bytes = checked_byte_extent(
        "swiglu output",
        &[total_rows, params.n as usize, DType::F32.size_of()],
    )?;
    if output.data_byte_len() < expected_out_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_swiglu_q4_0: output buffer too small: expected {} bytes, got {}",
            expected_out_bytes,
            output.data_byte_len()
        )));
    }
    let expected_ids_bytes =
        checked_byte_extent("swiglu ids", &[total_rows, DType::U32.size_of()])?;
    if ids.data_byte_len() < expected_ids_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "quantized_matmul_id_swiglu_q4_0: ids buffer too small: expected {} bytes, got {}",
            expected_ids_bytes,
            ids.data_byte_len(),
        )));
    }

    let pipeline =
        registry.get_pipeline("kernel_mul_mv_id_q4_0_f32_swiglu", device.metal_device())?;

    let gpu_params = GgmlMatvecIdGpuParams {
        ne00: params.k as i64,
        ne01: params.n as i64,
        ne02: params.n_experts as i64,
        ne10: params.k as i64,
        ne12: 1,
        ne0: params.n as i64,
        ne1: total_rows as i64,
        r2: 1,
        r3: 1,
        top_k: params.top_k,
        n_tokens: params.n_tokens,
        expert_stride: params.expert_stride as i64,
    };

    // Same (8, 8) geometry as the unfused Q4_0 mv_id kernel.
    let (nth0, nth1, align) = (8u64, 8u64, 8usize);
    let n = params.n as usize;
    let m = total_rows;
    let threadgroups = metal::MTLSize::new(div_ceil(n, align) as u64, m as u64, 1);
    let threads_per_tg = metal::MTLSize::new(nth0, nth1, 1);

    encoder.encode_threadgroups_with_args(
        pipeline,
        &[
            (0, KernelArg::Buffer(weight)),
            (1, KernelArg::Buffer(gate)),
            (2, KernelArg::Buffer(up)),
            (3, KernelArg::Buffer(output)),
            (4, KernelArg::Buffer(ids)),
            (5, KernelArg::Bytes(as_bytes(&gpu_params))),
        ],
        threadgroups,
        threads_per_tg,
    );

    Ok(())
}

/// Half-tile skip (P2c) selection for a compact `mm_id` launch.
#[doc(hidden)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HalfTileSkip {
    /// Enabled exactly on the kernel paths where it shipped (bitwise).
    Production,
    Disabled,
    /// Enabled on whichever path runs; parity and benchmark use only.
    Forced,
}

/// Complete `mm_id` schedule configuration. Production derives it from
/// [`GgmlRoutingPolicy::expert_mm_schedule`] via [`Self::production`];
/// parity tests and microbenchmarks pass explicit values.
///
/// The compact launch has one grid order, tile fastest: `(max_tiles,
/// ⌈N/64⌉, 1)`, every tile of one N block before the next. ADR-063 M0c
/// (Apple M5 Max, 2026-11-05, served shapes) measured it 0.012 s per
/// four-agent run faster than N tile fastest `(⌈N/64⌉, max_tiles, 1)`;
/// both give each output element a single writer, so the choice is
/// bitwise-neutral and N tile fastest was removed.
#[doc(hidden)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct IdMmScheduleOptions {
    pub schedule: ExpertMmSchedule,
    pub half_tile_skip: HalfTileSkip,
}

impl IdMmScheduleOptions {
    /// The configuration production dispatch uses for `schedule`.
    pub const fn production(schedule: ExpertMmSchedule) -> Self {
        Self {
            schedule,
            half_tile_skip: HalfTileSkip::Production,
        }
    }
}

/// P2c on the simdgroup path: simdgroups 2-3 own rows 16-31 and skip their
/// MMA and shared-memory store, which is bitwise by construction (§5.1 N5).
const HALF_TILE_SKIP_SIMDGROUP: bool = true;
/// P2c on the tensor-API path: a 16-row `matmul2d` on partial tiles. It is
/// not bitwise by construction (the 16-row product's reduction order is a
/// property of the device's `matmul2d` implementation), so it ships only on
/// the GPU family where §5.1 N2/N5 proved it bitwise: device names starting
/// with one of these prefixes (validated on "Apple M5 Max", 2026-11-04).
/// Other tensor-capable devices keep the 32-row body.
const HALF_TILE_SKIP_TENSOR_VALIDATED_DEVICES: &[&str] = &["Apple M5"];

/// Whether the tensor-path half-tile skip shipped on `device`. Cached per
/// Metal device so the hot path does not query the device name. Public for
/// parity and benchmark receipts.
#[doc(hidden)]
pub fn tensor_half_tile_skip_validated(device: &MlxDevice) -> bool {
    static CACHE: std::sync::Mutex<Option<(u64, bool)>> = std::sync::Mutex::new(None);
    let registry_id = device.registry_id();
    let mut cache = CACHE
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    match *cache {
        Some((id, validated)) if id == registry_id => validated,
        _ => {
            let name = device.name();
            let validated = HALF_TILE_SKIP_TENSOR_VALIDATED_DEVICES
                .iter()
                .any(|prefix| name.starts_with(prefix));
            *cache = Some((registry_id, validated));
            validated
        }
    }
}

const SCHED_FLAG_HALF_TILE_SKIP: u16 = 1;

/// `(htpe, hids)` bytes one schedule needs for `(n_experts, n_tokens, top_k)`.
fn id_mm_scratch_bytes(
    schedule: ExpertMmSchedule,
    n_experts: u32,
    n_tokens: u32,
    top_k: u32,
) -> Result<(usize, usize)> {
    let (htpe, hids) = expert_mm_schedule_scratch_bytes(schedule, n_experts, n_tokens, top_k)
        .ok_or_else(|| MlxError::InvalidArgument("mm_id schedule scratch size overflow".into()))?;
    let to_usize = |bytes: u64| {
        usize::try_from(bytes)
            .map_err(|_| MlxError::InvalidArgument("mm_id schedule scratch exceeds usize".into()))
    };
    Ok((to_usize(htpe)?, to_usize(hids)?))
}

/// Identity of the routing schedule a `map0` encode left in an
/// [`IdMmScratch`] (ADR-063 P2b).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct IdMmScheduleKey {
    ids_buffer: usize,
    ids_byte_offset: u64,
    ids_byte_extent: usize,
    n_tokens: u32,
    top_k: u32,
    n_experts: u32,
    schedule: ExpertMmSchedule,
    encoder: u64,
    encode_serial: u64,
}

/// Caller-owned scratch for the `_id` mm path's routing schedule.
///
/// Holds the two buffers `map0` writes and `mm_id` reads. Under
/// [`ExpertMmSchedule::UniformGrid`] they are `htpe` `[n_experts]`
/// per-expert counts and `hids` `[n_experts, n_tokens]` per-expert lists;
/// under [`ExpertMmSchedule::CompactTiles`] they are the tile descriptor
/// words (tile count plus three words per tile) and the compact
/// `rows + n_experts` list. Passing one instance through every mm_id call in
/// a prefill amortises what would otherwise be two Metal allocations per MoE
/// layer.
///
/// Size the scratch once for the largest call: [`Self::alloc`] bounds
/// `n_experts` and `n_tokens` per call and assumes `top_k <= 8` (the largest
/// `map0` instantiation); [`Self::alloc_with_top_k`] takes the top-k bound
/// explicitly. Either constructor serves both schedules, and `htpe` always
/// holds at least `n_experts` words for consumers of the legacy `map0`.
///
/// The buffers are private. Read them through [`Self::htpe`] and
/// [`Self::hids`]; every kernel that writes them must take them through
/// [`Self::htpe_mut`], [`Self::hids_mut`] or [`Self::buffers_mut`], which
/// invalidate the recorded schedule key that
/// [`quantized_matmul_id_ggml_pooled_slotted_reusing_schedule`] checks.
pub struct IdMmScratch {
    htpe: MlxBuffer,
    hids: MlxBuffer,
    n_experts_cap: u32,
    n_tokens_cap: u32,
    rows_cap: u32,
    encode_serial: u64,
    schedule_key: Option<IdMmScheduleKey>,
}

impl IdMmScratch {
    /// Allocate scratch for calls with at most `n_experts` experts and
    /// `max_n_tokens` tokens, assuming `top_k <= 8`.
    pub fn alloc(device: &MlxDevice, n_experts: u32, max_n_tokens: u32) -> Result<Self> {
        Self::alloc_with_top_k(device, n_experts, max_n_tokens, 8)
    }

    /// Allocate scratch for calls with at most `n_experts` experts,
    /// `max_n_tokens` tokens and `max_n_tokens * max_top_k` routed rows.
    pub fn alloc_with_top_k(
        device: &MlxDevice,
        n_experts: u32,
        max_n_tokens: u32,
        max_top_k: u32,
    ) -> Result<Self> {
        if n_experts == 0 || max_n_tokens == 0 || max_top_k == 0 {
            return Err(MlxError::InvalidArgument(
                "IdMmScratch::alloc dimensions must be nonzero".into(),
            ));
        }
        let rows_cap = max_n_tokens
            .checked_mul(max_top_k)
            .filter(|rows| *rows <= i32::MAX as u32)
            .ok_or_else(|| {
                MlxError::InvalidArgument(
                    "IdMmScratch: max_n_tokens * max_top_k exceeds signed Metal row indexing"
                        .into(),
                )
            })?;
        let (grid_htpe, grid_hids) =
            id_mm_scratch_bytes(ExpertMmSchedule::UniformGrid, n_experts, max_n_tokens, 1)?;
        let (compact_htpe, compact_hids) = id_mm_scratch_bytes(
            ExpertMmSchedule::CompactTiles,
            n_experts,
            max_n_tokens,
            max_top_k,
        )?;
        let htpe_bytes = grid_htpe.max(compact_htpe);
        let hids_bytes = grid_hids.max(compact_hids);
        let word = DType::U32.size_of();
        let htpe = device.alloc_buffer(htpe_bytes, DType::U32, vec![htpe_bytes / word])?;
        let hids = device.alloc_buffer(hids_bytes, DType::U32, vec![hids_bytes / word])?;
        Ok(Self {
            htpe,
            hids,
            n_experts_cap: n_experts,
            n_tokens_cap: max_n_tokens,
            rows_cap,
            encode_serial: 0,
            schedule_key: None,
        })
    }

    /// Largest `n_experts` a call may use.
    pub fn capacity_n_experts(&self) -> u32 {
        self.n_experts_cap
    }

    /// Largest `n_tokens` a call may use.
    pub fn capacity_n_tokens(&self) -> u32 {
        self.n_tokens_cap
    }

    /// Largest `n_tokens * top_k` a call may use.
    pub fn capacity_rows(&self) -> u32 {
        self.rows_cap
    }

    /// Check that a call fits the constructor's bounds.
    pub fn check_capacity(&self, n_experts: u32, n_tokens: u32, top_k: u32) -> Result<()> {
        if n_experts > self.n_experts_cap {
            return Err(MlxError::InvalidArgument(format!(
                "IdMmScratch: n_experts ({}) > cap ({})",
                n_experts, self.n_experts_cap,
            )));
        }
        if n_tokens > self.n_tokens_cap {
            return Err(MlxError::InvalidArgument(format!(
                "IdMmScratch: n_tokens ({}) > cap ({})",
                n_tokens, self.n_tokens_cap,
            )));
        }
        let rows = u64::from(n_tokens) * u64::from(top_k);
        if rows > u64::from(self.rows_cap) {
            return Err(MlxError::InvalidArgument(format!(
                "IdMmScratch: n_tokens * top_k ({rows}) > row cap ({})",
                self.rows_cap,
            )));
        }
        Ok(())
    }

    /// Schedule-word buffer, for host inspection or read-only binding.
    pub fn htpe(&self) -> &MlxBuffer {
        &self.htpe
    }

    /// Routed-row list buffer, for host inspection or read-only binding.
    pub fn hids(&self) -> &MlxBuffer {
        &self.hids
    }

    /// Mutable schedule-word buffer; invalidates the recorded schedule.
    pub fn htpe_mut(&mut self) -> &mut MlxBuffer {
        self.invalidate_schedule();
        &mut self.htpe
    }

    /// Mutable routed-row list buffer; invalidates the recorded schedule.
    pub fn hids_mut(&mut self) -> &mut MlxBuffer {
        self.invalidate_schedule();
        &mut self.hids
    }

    /// Both buffers for a kernel that writes the scratch (for example a
    /// legacy `map0` consumer); invalidates the recorded schedule.
    pub fn buffers_mut(&mut self) -> (&mut MlxBuffer, &mut MlxBuffer) {
        self.invalidate_schedule();
        (&mut self.htpe, &mut self.hids)
    }

    fn invalidate_schedule(&mut self) {
        self.encode_serial = self.encode_serial.wrapping_add(1);
        self.schedule_key = None;
    }

    fn schedule_key_for(
        &self,
        encoder: &CommandEncoder,
        ids: &MlxBuffer,
        n_tokens: u32,
        top_k: u32,
        n_experts: u32,
        schedule: ExpertMmSchedule,
    ) -> IdMmScheduleKey {
        IdMmScheduleKey {
            ids_buffer: &**ids.metal_buffer() as *const metal::BufferRef as usize,
            ids_byte_offset: ids.byte_offset(),
            ids_byte_extent: (n_tokens as usize)
                .saturating_mul(top_k as usize)
                .saturating_mul(DType::U32.size_of()),
            n_tokens,
            top_k,
            n_experts,
            schedule,
            encoder: encoder.instance_id(),
            encode_serial: self.encode_serial,
        }
    }

    fn validate_schedule_reuse(
        &self,
        encoder: &CommandEncoder,
        ids: &MlxBuffer,
        n_tokens: u32,
        top_k: u32,
        n_experts: u32,
        schedule: ExpertMmSchedule,
    ) -> Result<()> {
        let requested = self.schedule_key_for(encoder, ids, n_tokens, top_k, n_experts, schedule);
        match self.schedule_key {
            Some(recorded) if recorded == requested => Ok(()),
            Some(recorded) => Err(MlxError::InvalidArgument(format!(
                "IdMmScratch schedule reuse mismatch: recorded {recorded:?}, requested {requested:?}"
            ))),
            None => Err(MlxError::InvalidArgument(
                "IdMmScratch schedule reuse without a valid map0 schedule (none encoded, or invalidated by a mutable buffer access)".into(),
            )),
        }
    }
}

/// Matrix-matrix `_id` dispatch using caller-owned scratch (ADR-011 Phase
/// 3 Wave P3b).
#[allow(clippy::too_many_arguments)]
fn dispatch_id_mm_pooled(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    dispatch_id_mm_pooled_with_layout(
        encoder,
        registry,
        device,
        input,
        weight,
        ids,
        output,
        scratch,
        params,
        IdMmInputLayout::SharedPerToken,
        false,
        routing,
    )
}

#[allow(clippy::too_many_arguments)]
fn dispatch_id_mm_pooled_with_layout(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
    input_layout: IdMmInputLayout,
    schedule_prepared: bool,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    dispatch_id_mm_pooled_with_options(
        encoder,
        registry,
        device,
        input,
        weight,
        ids,
        output,
        scratch,
        params,
        input_layout,
        schedule_prepared,
        routing,
        IdMmScheduleOptions::production(routing.expert_mm_schedule),
    )
}

#[allow(clippy::too_many_arguments)]
fn dispatch_id_mm_pooled_with_options(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
    input_layout: IdMmInputLayout,
    schedule_prepared: bool,
    routing: &GgmlRoutingPolicy,
    options: IdMmScheduleOptions,
) -> Result<()> {
    scratch.check_capacity(params.n_experts, params.n_tokens, params.top_k)?;

    // Translate the dispatcher-facing params to the mm_id internal dispatch
    // shape.  Same fields; different type keeps the public mv params from
    // becoming mm-specific.
    let dispatch = GgmlIdMmDispatchParams {
        n_tokens: params.n_tokens,
        top_k: params.top_k,
        n: params.n,
        k: params.k,
        n_experts: params.n_experts,
        expert_stride: params.expert_stride,
        ggml_type: params.ggml_type,
    };

    let result = dispatch_id_mm_with_layout(
        encoder,
        registry,
        device,
        input,
        weight,
        ids,
        &scratch.htpe,
        &scratch.hids,
        output,
        &dispatch,
        input_layout,
        schedule_prepared,
        routing,
        options,
    );
    match (&result, schedule_prepared) {
        // A fresh schedule is now the scratch's schedule.
        (Ok(()), false) => {
            scratch.encode_serial = scratch.encode_serial.wrapping_add(1);
            scratch.schedule_key = Some(scratch.schedule_key_for(
                encoder,
                ids,
                params.n_tokens,
                params.top_k,
                params.n_experts,
                options.schedule,
            ));
        }
        (Ok(()), true) => {}
        // Fail closed: a failed call may have encoded part of a schedule.
        (Err(_), _) => scratch.invalidate_schedule(),
    }
    result
}

/// Test-only helper: pooled `mm_id` under an explicit schedule
/// configuration (schedule, half-tile skip), for parity tests and the
/// schedule microbenchmark.
///
/// `slotted` selects the `[n_tokens, top_k, K]` input layout. With
/// `reuse_schedule` the call validates the scratch's recorded schedule key
/// exactly like [`quantized_matmul_id_ggml_pooled_slotted_reusing_schedule`]
/// and encodes no `map0`. The call must be on the `mm_id` route.
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
pub fn dispatch_id_mm_pooled_with_options_for_test(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    scratch: &mut IdMmScratch,
    params: &GgmlQuantizedMatmulIdParams,
    slotted: bool,
    reuse_schedule: bool,
    routing: &GgmlRoutingPolicy,
    options: IdMmScheduleOptions,
) -> Result<()> {
    validate_native_expert_dtypes(
        "dispatch_id_mm_pooled_with_options_for_test",
        input,
        weight,
        ids,
        output,
    )?;
    if plan_expert_auto_route(params.n_tokens, params.top_k, params.k, false, routing)
        != ExpertAutoPlan::Mm
    {
        return Err(MlxError::InvalidArgument(
            "dispatch_id_mm_pooled_with_options_for_test requires the mm_id route".into(),
        ));
    }
    if reuse_schedule {
        scratch.validate_schedule_reuse(
            encoder,
            ids,
            params.n_tokens,
            params.top_k,
            params.n_experts,
            options.schedule,
        )?;
    }
    dispatch_id_mm_pooled_with_options(
        encoder,
        registry,
        device,
        input,
        weight,
        ids,
        output,
        scratch,
        params,
        if slotted {
            IdMmInputLayout::Slotted
        } else {
            IdMmInputLayout::SharedPerToken
        },
        reuse_schedule,
        routing,
        options,
    )
}

/// Matrix-matrix `_id` dispatch that allocates scratch on every call.
///
/// Retained for the auto-allocating `quantized_matmul_id_ggml` entry
/// point (tests, non-prefill callers); the pooled entry point
/// `quantized_matmul_id_ggml_pooled` is preferred for batched prefill.
#[allow(clippy::too_many_arguments)]
fn dispatch_id_mm(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlQuantizedMatmulIdParams,
    routing: &GgmlRoutingPolicy,
) -> Result<()> {
    let mut scratch =
        IdMmScratch::alloc_with_top_k(device, params.n_experts, params.n_tokens, params.top_k)?;
    dispatch_id_mm_pooled(
        encoder,
        registry,
        device,
        input,
        weight,
        ids,
        output,
        &mut scratch,
        params,
        routing,
    )
}

fn div_ceil(a: usize, b: usize) -> usize {
    (a + b - 1) / b
}

// ============================================================================
// ADR-011 Phase 3 Wave P3a: `_id` matrix-matrix (mm) path.
//
// Peer port of the `kernel_mul_mm_id_map0_ne20_<N>` + `kernel_mul_mm_id_<q>_f32`
// two-stage dispatch.  Used for MoE projections at prefill — instead of
// re-reading each expert's weight blocks once per routed (token, slot) pair,
// the mm kernel stages a 64x32 expert weight tile into threadgroup shared
// memory and reuses it across a 32-row block of the expert's routed tokens.
//
// The preprocessing step (`map0`) is what lets mm work for MoE: it
// regroups the flat `[n_tokens, top_k]` ids table into per-expert routed
// token lists so each mm tile is homogeneous in its choice of expert
// weight slab.  Without map0, consecutive M-rows in a tile could route to
// different experts, defeating weight reuse.
//
// Same staging strategy as Wave P3a Commit 1 (non-id): the kernel exists,
// tests verify correctness, but the public `quantized_matmul_id_ggml`
// dispatcher is NOT rerouted yet — tests call `dispatch_id_mm_for_test`.
// ============================================================================

/// Host-side params for the `_id` mm path's `map0` preprocessor.
///
/// Matches `GgmlMatmulIdMm_Map0Params` in
/// `/opt/mlx-native/src/shaders/quantized_matmul_id_mm.metal`.  Explicit
/// 4-byte trailing padding so the struct aligns to 8 (u64).
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct GgmlIdMmMap0GpuParams {
    ne10: i32, // unused, kept for struct symmetry
    ne11: i32, // n_expert_used (bcast, == ne20)
    nb11: u64, // unused
    nb12: u64, // unused
    ne21: i32, // n_tokens
    ne20: i32, // n_expert_used (top_k)
    nb21: u64, // bytes per token in the ids table (= ne20 * sizeof(i32))
}

/// Host-side params for the `_id` mm kernel.
///
/// Matches `GgmlMatmulIdMm_MmParams` in
/// `/opt/mlx-native/src/shaders/quantized_matmul_id_mm.metal`.
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct GgmlIdMmMmGpuParams {
    ne00: i32, // K
    ne02: i32, // n_experts
    nb01: u64, // bytes per weight row (within one expert's slab)
    nb02: u64, // bytes per expert weight slab (= nb01 * N)
    nb03: u64,
    ne11: i32, // n_expert_used (bcast)
    _pad0: u32,
    nb10: u64, // = sizeof(float)
    nb11: u64, // bytes per input row (= K * 4)
    nb12: u64, // bytes per input batch (= n_tokens * nb11)
    nb13: u64,
    ne20: i32, // n_expert_used (top_k)
    ne21: i32, // n_tokens
    ne0: i32,  // N (per-expert output rows)
    ne1: i32,  // batch stride (== ne20 for our packed layout)
    r2: i16,
    r3: i16,
    /// Compact kernels only: bit 0 half-tile skip.
    sched_flags: u16,
    _pad1: u16,
}

/// Parameters for the `_id` mm dispatch (scratch-buffer sized view).
#[derive(Debug, Clone, Copy)]
pub struct GgmlIdMmDispatchParams {
    /// Number of input tokens.
    pub n_tokens: u32,
    /// Number of experts each token is routed to (top-k).
    pub top_k: u32,
    /// Number of output columns per expert (weight rows).
    pub n: u32,
    /// Input dimension (weight cols before quantization).
    pub k: u32,
    /// Total experts in the stacked weight buffer.
    pub n_experts: u32,
    /// Byte stride between expert weight slices in the stacked buffer.
    pub expert_stride: u64,
    /// GGML quantization type.
    pub ggml_type: GgmlType,
}

impl GgmlIdMmDispatchParams {
    /// Bytes required for the `htpe` scratch buffer (per-expert routed count).
    pub fn htpe_bytes(&self) -> usize {
        (self.n_experts as usize).saturating_mul(DType::U32.size_of())
    }

    /// Bytes required for the `hids` scratch buffer (per-expert routed-token list).
    /// Layout: `[n_experts, n_tokens]` int32 row-major.
    ///
    /// Per-expert routed count is bounded by `n_tokens` when every token's
    /// top-k IDs are distinct. Map0 validates that condition on-device before
    /// appending; duplicate or out-of-range metadata poisons the schedule and
    /// is never allowed to exceed this fixed-capacity list.
    pub fn hids_bytes(&self) -> usize {
        (self.n_experts as usize)
            .saturating_mul(self.n_tokens as usize)
            .saturating_mul(DType::U32.size_of())
    }
}

/// Test-only helper: force the `_id` mm two-stage dispatch path.
///
/// Runs `kernel_mul_mm_id_map0_ne20_<top_k>` followed by
/// `kernel_mul_mm_id_<qtype>_f32`.
///
/// Input:
///   * `input`   — f32 input rows `[n_tokens, K]`.
///   * `weight`  — stacked expert weights `[n_experts, N, packed_K]`.
///   * `ids`     — flat expert-id table `[n_tokens, top_k]` viewed as
///                 i32 (u32 is byte-equivalent in this range).
///   * `output`  — f32 output `[n_tokens, top_k, N]` row-major.
///
/// Scratch (caller-allocated, zero-init not required; written by `map0`, so
/// taken mutably — an [`IdMmScratch`] lends them only through its
/// invalidating [`IdMmScratch::buffers_mut`]):
///   * `htpe`    — `[n_experts]` u32 (per-expert count).
///   * `hids`    — `[n_experts, n_tokens]` i32 (per-expert routed list).
///
/// Not intended for production callers — the public `quantized_matmul_id_ggml`
/// entry point stays on the mv path until the follow-up commit wires the
/// m > 8 threshold.
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
pub fn dispatch_id_mm_for_test(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    htpe: &mut MlxBuffer,
    hids: &mut MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlIdMmDispatchParams,
) -> Result<()> {
    let routing = GgmlRoutingPolicy::default();
    dispatch_id_mm_with_layout(
        encoder,
        registry,
        device,
        input,
        weight,
        ids,
        &*htpe,
        &*hids,
        output,
        params,
        IdMmInputLayout::SharedPerToken,
        false,
        &routing,
        IdMmScheduleOptions::production(ExpertMmSchedule::UniformGrid),
    )
}

/// Test-only helper: run `_id` matrix-matrix using a routing schedule that a
/// preceding [`dispatch_id_mm_for_test`] call prepared in the same scratch
/// buffers.
///
/// This isolates the large-prefill schedule-reuse hypothesis without changing
/// production routing. The preceding call and this call must use the exact
/// same `ids`, `n_tokens`, `top_k`, `n_experts`, `htpe`, and `hids`. They must
/// be encoded into the same command encoder (the first call's map0 barrier
/// orders schedule construction), or the first encoder must have completed.
/// Gate and up projections may differ in weight buffer and expert stride but
/// must otherwise have compatible projection shapes.
///
/// Violating that contract can read a stale routing schedule. Production
/// callers must use [`quantized_matmul_id_ggml_pooled_pair`], which owns both
/// projections and does not expose schedule identity or lifetime.
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
pub fn dispatch_id_mm_with_prepared_schedule_for_test(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    htpe: &MlxBuffer,
    hids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlIdMmDispatchParams,
) -> Result<()> {
    let routing = GgmlRoutingPolicy::default();
    dispatch_id_mm_with_layout(
        encoder,
        registry,
        device,
        input,
        weight,
        ids,
        htpe,
        hids,
        output,
        params,
        IdMmInputLayout::SharedPerToken,
        true,
        &routing,
        IdMmScheduleOptions::production(ExpertMmSchedule::UniformGrid),
    )
}

/// Test-only helper: encode only the routing-schedule stage of `schedule`.
///
/// `UniformGrid` runs the legacy `map0` (`htpe` `[n_experts]` counts, `hids`
/// `[n_experts, n_tokens]` lists); `CompactTiles` runs `map0c` (tile count,
/// descriptors and the compact list). No matrix kernel is encoded.
/// Microbenchmarks use it to time the schedule stage on its own, and
/// parity tests read the schedule back on the host.
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
pub fn dispatch_id_mm_map0_for_test(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    ids: &MlxBuffer,
    htpe: &mut MlxBuffer,
    hids: &mut MlxBuffer,
    params: &GgmlIdMmDispatchParams,
    schedule: ExpertMmSchedule,
) -> Result<()> {
    let (htpe, hids) = (&*htpe, &*hids);
    validate_id_mm_scratch_dtypes("dispatch_id_mm_map0_for_test", htpe, hids)?;
    if ids.dtype() != DType::U32 {
        return Err(MlxError::InvalidArgument(format!(
            "dispatch_id_mm_map0_for_test requires U32 expert IDs, got {:?}",
            ids.dtype()
        )));
    }
    if params.n_tokens == 0
        || params.top_k == 0
        || params.n_experts == 0
        || params.top_k > params.n_experts
        || params.n_tokens > i32::MAX as u32
    {
        return Err(MlxError::InvalidArgument(
            "dispatch_id_mm_map0_for_test: n_tokens, top_k and n_experts must be valid".into(),
        ));
    }
    let total_rows = (params.n_tokens as usize)
        .checked_mul(params.top_k as usize)
        .ok_or_else(|| MlxError::InvalidArgument("map0 row count overflow".into()))?;
    let ids_bytes = checked_byte_extent("map0 ids", &[total_rows, DType::U32.size_of()])?;
    let (htpe_bytes, hids_bytes) =
        id_mm_scratch_bytes(schedule, params.n_experts, params.n_tokens, params.top_k)?;
    if ids.data_byte_len() < ids_bytes
        || htpe.data_byte_len() < htpe_bytes
        || hids.data_byte_len() < hids_bytes
    {
        return Err(MlxError::InvalidArgument(
            "dispatch_id_mm_map0_for_test: ids or scratch buffer too small".into(),
        ));
    }
    match schedule {
        ExpertMmSchedule::UniformGrid => {
            encode_legacy_map0(encoder, registry, device, ids, htpe, hids, params)
        }
        ExpertMmSchedule::CompactTiles => encode_compact_map0(
            encoder,
            registry,
            device,
            ids,
            htpe,
            hids,
            params,
            compact_max_tiles(total_rows as u64, u64::from(params.n_experts)),
        ),
    }
}

/// Encode the legacy `kernel_mul_mm_id_map0_ne20_<top_k>` stage plus the
/// barrier that orders it before the consuming `mm_id` kernel.
fn encode_legacy_map0(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    ids: &MlxBuffer,
    htpe: &MlxBuffer,
    hids: &MlxBuffer,
    params: &GgmlIdMmDispatchParams,
) -> Result<()> {
    let map0_kernel_name = match params.top_k {
        1 => "kernel_mul_mm_id_map0_ne20_1",
        6 => "kernel_mul_mm_id_map0_ne20_6",
        8 => "kernel_mul_mm_id_map0_ne20_8",
        other => {
            return Err(MlxError::InvalidArgument(format!(
                "dispatch_id_mm_for_test: no map0 instantiation for top_k={}",
                other
            )))
        }
    };
    let map0_pipeline = registry.get_pipeline(map0_kernel_name, device.metal_device())?;
    if u64::from(params.n_experts) > map0_pipeline.max_total_threads_per_threadgroup() {
        return Err(MlxError::InvalidArgument(format!(
            "n_experts ({}) exceeds map0 pipeline threadgroup limit ({})",
            params.n_experts,
            map0_pipeline.max_total_threads_per_threadgroup(),
        )));
    }

    let map0_params = GgmlIdMmMap0GpuParams {
        ne10: params
            .n
            .try_into()
            .map_err(|_| MlxError::InvalidArgument("N out of i32 range".into()))?,
        ne11: params.top_k as i32,
        nb11: 0,
        nb12: 0,
        ne21: params.n_tokens as i32,
        ne20: params.top_k as i32,
        nb21: (params.top_k as u64) * (DType::U32.size_of() as u64),
    };

    let map0_shmem =
        (params.n_experts as u64) * (params.top_k as u64) * DType::U32.size_of() as u64;
    let map0_threadgroups = metal::MTLSize::new(1, 1, 1);
    let map0_threads = metal::MTLSize::new(params.n_experts as u64, 1, 1);

    if encoder.is_capturing() {
        let range = |buffer: &MlxBuffer| {
            let start =
                (buffer.contents_ptr() as usize).saturating_add(buffer.byte_offset() as usize);
            (start, start.saturating_add(buffer.data_byte_len()))
        };
        encoder.set_pending_buffer_ranges(vec![range(ids)], vec![range(htpe), range(hids)]);
    }
    encoder.encode_threadgroups_with_args_and_shared(
        map0_pipeline,
        &[
            (0, KernelArg::Bytes(as_bytes(&map0_params))),
            (1, KernelArg::Buffer(ids)),
            (2, KernelArg::Buffer(htpe)),
            (3, KernelArg::Buffer(hids)),
        ],
        &[(0, map0_shmem)],
        map0_threadgroups,
        map0_threads,
    );

    // Memory barrier: the mm kernel reads htpe + hids, map0 wrote them.
    // Without this, Metal's concurrent-dispatch compute encoder lets the
    // two dispatches overlap — mm would read zeros (all-expert early-exit).
    encoder.memory_barrier();
    Ok(())
}

/// Host-side params for `kernel_mul_mm_id_map0c_ne20_<N>`; matches
/// `GgmlMatmulIdMm_Map0cParams` in `quantized_matmul_id_mm.metal`.
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct GgmlIdMmMap0cGpuParams {
    ne21: i32,      // n_tokens
    ne20: i32,      // top_k
    nb21: u64,      // bytes per token in the ids table
    max_tiles: u32, // descriptor capacity; more tiles poison the schedule
    n_experts: u32,
}

/// Encode the compact `kernel_mul_mm_id_map0c_ne20_<top_k>` stage (ADR-063
/// P2) plus the barrier that orders it before the consuming `mm_id` kernel.
#[allow(clippy::too_many_arguments)]
fn encode_compact_map0(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    ids: &MlxBuffer,
    htpe: &MlxBuffer,
    hids: &MlxBuffer,
    params: &GgmlIdMmDispatchParams,
    max_tiles: u64,
) -> Result<()> {
    let kernel_name = match params.top_k {
        1 => "kernel_mul_mm_id_map0c_ne20_1",
        6 => "kernel_mul_mm_id_map0c_ne20_6",
        8 => "kernel_mul_mm_id_map0c_ne20_8",
        other => {
            return Err(MlxError::InvalidArgument(format!(
                "dispatch_id_mm: no map0c instantiation for top_k={other}"
            )))
        }
    };
    let pipeline = registry.get_pipeline(kernel_name, device.metal_device())?;
    // Token-parallel launch of 32·g threads, g ≤ 32 simdgroups: one per 32
    // tokens (or experts) up to the pipeline and threadgroup-memory limits.
    // Threadgroup memory: per-expert positions, max(g, 2) lane masks per
    // expert, 2·g scan partials and the invalid flag, in u32 words.
    let n_experts = u64::from(params.n_experts);
    let word = DType::U32.size_of() as u64;
    let shmem_bytes =
        |g: u64| ((n_experts + g.max(2) * n_experts + 2 * g + 1) * word).div_ceil(16) * 16;
    let shmem_budget = device
        .metal_device()
        .max_threadgroup_memory_length()
        .saturating_sub(pipeline.static_threadgroup_memory_length());
    let mut simdgroups = u64::from(params.n_tokens)
        .max(n_experts)
        .div_ceil(32)
        .min(32)
        .min(pipeline.max_total_threads_per_threadgroup() / 32);
    while simdgroups > 0 && shmem_bytes(simdgroups) > shmem_budget {
        simdgroups -= 1;
    }
    if simdgroups == 0 {
        return Err(MlxError::InvalidArgument(format!(
            "n_experts ({}) exceeds the map0c threadgroup memory ({shmem_budget} bytes)",
            params.n_experts,
        )));
    }
    let map0_params = GgmlIdMmMap0cGpuParams {
        ne21: params.n_tokens as i32,
        ne20: params.top_k as i32,
        nb21: u64::from(params.top_k) * DType::U32.size_of() as u64,
        max_tiles: u32::try_from(max_tiles)
            .map_err(|_| MlxError::InvalidArgument("compact tile bound exceeds u32".into()))?,
        n_experts: params.n_experts,
    };
    let shmem = shmem_bytes(simdgroups);

    if encoder.is_capturing() {
        let range = |buffer: &MlxBuffer| {
            let start =
                (buffer.contents_ptr() as usize).saturating_add(buffer.byte_offset() as usize);
            (start, start.saturating_add(buffer.data_byte_len()))
        };
        encoder.set_pending_buffer_ranges(vec![range(ids)], vec![range(htpe), range(hids)]);
    }
    encoder.encode_threadgroups_with_args_and_shared(
        pipeline,
        &[
            (0, KernelArg::Bytes(as_bytes(&map0_params))),
            (1, KernelArg::Buffer(ids)),
            (2, KernelArg::Buffer(htpe)),
            (3, KernelArg::Buffer(hids)),
        ],
        &[(0, shmem)],
        metal::MTLSize::new(1, 1, 1),
        metal::MTLSize::new(32 * simdgroups, 1, 1),
    );
    // The compact mm_id kernel reads the descriptors and list map0c wrote.
    encoder.memory_barrier();
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn dispatch_id_mm_with_layout(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    ids: &MlxBuffer,
    htpe: &MlxBuffer,
    hids: &MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlIdMmDispatchParams,
    input_layout: IdMmInputLayout,
    schedule_prepared: bool,
    routing: &GgmlRoutingPolicy,
    options: IdMmScheduleOptions,
) -> Result<()> {
    validate_native_expert_dtypes("dispatch_id_mm", input, weight, ids, output)?;
    validate_id_mm_scratch_dtypes("dispatch_id_mm", htpe, hids)?;
    let qk = params.ggml_type.block_values();

    // ---- Validate common shapes ----
    // ADR-013 P16 — Q4_K added. ADR-022 — Q5_1 / IQ4_NL added.
    // ADR-022 Phase 2 — Q5_K added. ADR-033 §Pi Task #20 — IQ4_XS added.
    // ADR-066 N-IQ2XXS — IQ2_XXS added.
    match params.ggml_type {
        GgmlType::Q4_0
        | GgmlType::Q5_0
        | GgmlType::Q8_0
        | GgmlType::Q2_K
        | GgmlType::Q3_K
        | GgmlType::Q4_K
        | GgmlType::Q5_K
        | GgmlType::Q6_K
        | GgmlType::Q5_1
        | GgmlType::IQ4_NL
        | GgmlType::IQ4_XS
        | GgmlType::IQ2_XXS => {}
        other => {
            return Err(MlxError::InvalidArgument(format!(
                "dispatch_id_mm_for_test does not support {:?}",
                other
            )));
        }
    }
    if params.n_tokens == 0
        || params.k == 0
        || params.n == 0
        || params.top_k == 0
        || params.n_experts == 0
    {
        return Err(MlxError::InvalidArgument(
            "n_tokens, K, N, top_k, n_experts must all be > 0".into(),
        ));
    }
    if params.top_k > params.n_experts {
        return Err(MlxError::InvalidArgument(
            "top_k must not exceed n_experts".into(),
        ));
    }
    for (label, value) in [
        ("n_tokens", params.n_tokens),
        ("K", params.k),
        ("N", params.n),
        ("top_k", params.top_k),
        ("n_experts", params.n_experts),
    ] {
        if value > i32::MAX as u32 {
            return Err(MlxError::InvalidArgument(format!(
                "dispatch_id_mm_for_test: {label} exceeds the signed Metal ABI"
            )));
        }
    }
    if params.k % qk != 0 {
        return Err(MlxError::InvalidArgument(format!(
            "K ({}) must be divisible by block QK ({})",
            params.k, qk
        )));
    }

    // Match the top_k template instantiations available in the shader.
    // P3b-tensor.2 — added ne20_1 alongside ne20_8 (Gemma 4 MoE down).
    if !has_mm_id_map0(params.top_k) {
        return Err(MlxError::InvalidArgument(format!(
            "dispatch_id_mm_for_test: top_k {} has no map0 instantiation (need 1, 6, or 8)",
            params.top_k
        )));
    }

    let blocks_per_row = params.k / qk;
    let block_bytes = params.ggml_type.block_bytes();
    let total_weight_bytes = required_expert_weight_bytes(
        params.ggml_type,
        params.n_experts,
        params.n,
        params.k,
        params.expert_stride,
    )?;
    if weight.data_byte_len() < total_weight_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "dispatch_id_mm_for_test: weight buffer too small: expected {} bytes, got {}",
            total_weight_bytes,
            weight.data_byte_len(),
        )));
    }
    let input_rows = match input_layout {
        IdMmInputLayout::SharedPerToken => params.n_tokens as usize,
        IdMmInputLayout::Slotted => (params.n_tokens as usize)
            .checked_mul(params.top_k as usize)
            .ok_or_else(|| MlxError::InvalidArgument("slotted input row count overflow".into()))?,
    };
    let input_bytes = checked_byte_extent(
        "mm input",
        &[input_rows, params.k as usize, DType::F32.size_of()],
    )?;
    if input.data_byte_len() < input_bytes {
        return Err(MlxError::InvalidArgument(
            "dispatch_id_mm_for_test: input buffer too small".into(),
        ));
    }
    let total_rows = (params.n_tokens as usize)
        .checked_mul(params.top_k as usize)
        .ok_or_else(|| MlxError::InvalidArgument("mm expert row count overflow".into()))?;
    if total_rows > i32::MAX as usize {
        return Err(MlxError::InvalidArgument(
            "dispatch_id_mm_for_test: n_tokens * top_k exceeds signed Metal row indexing".into(),
        ));
    }
    let ids_bytes = checked_byte_extent("mm ids", &[total_rows, DType::U32.size_of()])?;
    if ids.data_byte_len() < ids_bytes {
        return Err(MlxError::InvalidArgument(
            "dispatch_id_mm_for_test: ids buffer too small".into(),
        ));
    }
    let output_bytes = checked_byte_extent(
        "mm output",
        &[total_rows, params.n as usize, DType::F32.size_of()],
    )?;
    if output.data_byte_len() < output_bytes {
        return Err(MlxError::InvalidArgument(
            "dispatch_id_mm_for_test: output buffer too small".into(),
        ));
    }
    let (htpe_bytes, hids_bytes) = id_mm_scratch_bytes(
        options.schedule,
        params.n_experts,
        params.n_tokens,
        params.top_k,
    )?;
    if htpe.data_byte_len() < htpe_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "dispatch_id_mm_for_test: htpe buffer too small for the {:?} schedule: expected {} bytes, got {}",
            options.schedule,
            htpe_bytes,
            htpe.data_byte_len(),
        )));
    }
    if hids.data_byte_len() < hids_bytes {
        return Err(MlxError::InvalidArgument(format!(
            "dispatch_id_mm_for_test: hids buffer too small for the {:?} schedule: expected {} bytes, got {}",
            options.schedule,
            hids_bytes,
            hids.data_byte_len(),
        )));
    }
    let compact = options.schedule == ExpertMmSchedule::CompactTiles;
    let max_tiles = compact_max_tiles(total_rows as u64, u64::from(params.n_experts));

    // ---- Stage 1: map0 — build per-expert routed-token lists ----
    //
    // Dispatch: legacy map0 runs 1 threadgroup of `n_experts` threads with a
    // `n_experts * top_k * sizeof(uint32)` staging area; compact map0c runs one
    // token-parallel threadgroup (`encode_compact_map0`). Both validate ids
    // (bounds, per-token distinctness) before any id indexes memory.
    //
    // ADR-011 Phase 3 Wave P3b-tensor.2 — pick the map0 instantiation
    // whose `ne20` template arg matches our top_k.  Gemma 4 needs both:
    // `ne20_8` for the gate_up call (top_k=8) and `ne20_1` for the
    // down call (top_k=1, where each output row routes to a single
    // expert).  Without ne20_1 the top_k=1 caller falls back to mv_id
    // and re-reads each expert's weights once per (seq_len*top_k) row —
    // ~50% of prefill time wasted on weight re-reads.
    if !schedule_prepared {
        if compact {
            encode_compact_map0(
                encoder, registry, device, ids, htpe, hids, params, max_tiles,
            )?;
        } else {
            encode_legacy_map0(encoder, registry, device, ids, htpe, hids, params)?;
        }
    }

    // ---- Stage 2: mm_id — matmul with the per-expert lists ----
    //
    // ADR-011 Phase 3 Wave P3b-tensor — prefer the tensor_ops::matmul2d
    // mm_id variant on M3+. The device-bound registry caches the exact
    // pipeline result after the first dispatch.
    let use_tensor = tensor_mm_auto_selected(
        params.ggml_type,
        routing.expert_tensor_mm,
        routing.q5_0_tensor_mm,
    ) && probe_tensor_mm_id(registry, device)?;
    // Fall through to the simdgroup variant when the per-type tensor
    // kernel isn't shipped — e.g. IQ4_XS (ADR-033 §Pi Task #20 shipped
    // the simdgroup variant only; tensor-API variant deferred as a
    // perf optimization).
    let tensor_name = if compact {
        params.ggml_type.id_mm_tensor_compact_kernel_name()
    } else {
        params.ggml_type.id_mm_tensor_kernel_name()
    };
    let tensor_path = use_tensor && tensor_name != "unsupported";
    let mm_kernel_name = if tensor_path {
        tensor_name
    } else if compact {
        params.ggml_type.id_mm_compact_kernel_name()
    } else {
        params.ggml_type.id_mm_kernel_name()
    };
    let half_tile_skip = compact
        && match options.half_tile_skip {
            HalfTileSkip::Production => {
                if tensor_path {
                    tensor_half_tile_skip_validated(device)
                } else {
                    HALF_TILE_SKIP_SIMDGROUP
                }
            }
            HalfTileSkip::Disabled => false,
            HalfTileSkip::Forced => true,
        };
    let sched_flags = if half_tile_skip {
        SCHED_FLAG_HALF_TILE_SKIP
    } else {
        0
    };
    let mm_pipeline = registry.get_pipeline(mm_kernel_name, device.metal_device())?;

    let nb01 = (blocks_per_row as u64) * (block_bytes as u64);
    let row_bytes = (params.k as u64) * (DType::F32.size_of() as u64);
    let (slot_stride, token_stride) = match input_layout {
        IdMmInputLayout::SharedPerToken => (0, row_bytes),
        IdMmInputLayout::Slotted => (
            row_bytes,
            row_bytes
                .checked_mul(params.top_k as u64)
                .ok_or_else(|| MlxError::InvalidArgument("slotted token stride overflow".into()))?,
        ),
    };

    // The normal input layout is `[n_tokens, K]`: nb11=0 shares one input
    // row across every expert slot and nb12=K*4 advances by token. Slotted
    // MoE down inputs are `[n_tokens, top_k, K]`: nb11=K*4 advances by slot
    // and nb12=top_k*K*4 advances by token. The latter matches the
    // reference MUL_MAT_ID addressing without forcing callers to flatten
    // top_k into the token count.
    let mm_params = GgmlIdMmMmGpuParams {
        ne00: params.k as i32,
        ne02: params.n_experts as i32,
        nb01,
        nb02: params.expert_stride,
        nb03: 0,
        ne11: params.top_k as i32,
        _pad0: 0,
        nb10: DType::F32.size_of() as u64,
        nb11: slot_stride,
        nb12: token_stride,
        nb13: 0,
        ne20: params.top_k as i32,
        ne21: params.n_tokens as i32,
        ne0: params.n as i32,
        ne1: params.top_k as i32,
        r2: 1,
        r3: 1,
        sched_flags,
        _pad1: 0,
    };

    const NR0: u64 = 64;
    const NR1: u64 = 32;
    const THREADS_PER_TG: u64 = 128;

    // Grid X: per-expert routed-row tiles. Kernel uses `r1 = tgpig.x *
    // NR1` as a routed-row index within an expert's hids list. Each
    // expert's hids list has ≤ n_tokens entries (production MoE top-k
    // routing selects distinct experts per token — each token
    // contributes ≤ 1 to any single expert). The kernel's early-exit
    // at line 534 (`if (r1 >= neh1) return`) drops tiles for experts
    // with fewer routed tokens.
    //
    // map0 validates the former caller invariant on-device before appending:
    // every ID must be in range and the top-k IDs for one token must be
    // distinct. Invalid metadata sets a poison bit in every htpe entry; the
    // consumer then writes NaN to the entire output without reading weights.
    // This stays on-device because production routers can create IDs earlier
    // in this same command stream; a host scan would observe stale data or
    // introduce a forbidden synchronization point.
    //
    // Compact schedule: map0c emitted at most `max_tiles` descriptors and
    // stored the live count in word 0; tile slots at or past it return.
    let n_tiles = (params.n as u64 + NR0 - 1) / NR0;
    let mm_threadgroups = if !compact {
        metal::MTLSize::new(
            (params.n_tokens as u64 + NR1 - 1) / NR1,
            n_tiles,
            params.n_experts as u64,
        )
    } else {
        metal::MTLSize::new(max_tiles, n_tiles, 1)
    };
    let mm_threads = metal::MTLSize::new(THREADS_PER_TG, 1, 1);

    const MM_SHMEM_BYTES: u64 = 8192;

    if encoder.is_capturing() {
        let range = |buffer: &MlxBuffer| {
            let start =
                (buffer.contents_ptr() as usize).saturating_add(buffer.byte_offset() as usize);
            (start, start.saturating_add(buffer.data_byte_len()))
        };
        encoder.set_pending_buffer_ranges(
            vec![range(weight), range(input), range(htpe), range(hids)],
            vec![range(output)],
        );
    }
    encoder.encode_threadgroups_with_args_and_shared(
        mm_pipeline,
        &[
            (0, KernelArg::Bytes(as_bytes(&mm_params))),
            (1, KernelArg::Buffer(weight)),
            (2, KernelArg::Buffer(input)),
            (3, KernelArg::Buffer(htpe)),
            (4, KernelArg::Buffer(hids)),
            (5, KernelArg::Buffer(output)),
        ],
        &[(0, MM_SHMEM_BYTES)],
        mm_threadgroups,
        mm_threads,
    );

    Ok(())
}

// ============================================================================
// ADR-033 §Pi Task #20 + ADR-034 §93 — fused MoE gate_up_silu_mm_id dispatch.
//
// Replaces 3 dispatches per MoE FFN per layer:
//
//   1. gate_mm_id  — matmul gate_w × x  → tmp_gate
//   2. up_mm_id    — matmul up_w   × x  → tmp_up
//   3. silu_mul_id — elem-wise   silu(tmp_gate) * tmp_up → out
//
// with a single fused dispatch:
//
//   fused_gate_up_silu_mm_id(gate_w, up_w, x, htpe, hids, out)
//
// The map0 stage runs once (shared between gate and up — same ids), then
// the fused kernel reuses the staged B-tile across both gate and up matmuls
// and applies silu_mul element-wise at writeback. Bandwidth savings:
//   - 1× input B-tile read instead of 2×
//   - 1× output writeback instead of 2× (tmp_gate + tmp_up no longer needed)
//   - 1× map0 dispatch instead of 2× (shared in unfused path)
//   - Dispatch overhead: 1 kernel launch instead of 3
//
// Currently Q6_K only — Q5_K and IQ4_XS variants are next iters of the
// same arc per ADR-033 §Pi multi-iter plan.
// ============================================================================

/// Test-only helper: invoke the fused gate_up_silu mm_id kernel.
///
/// Computes `out[token*top_k + slot][col] = silu(gate_w[expert][col] · x[token])
///                                          * (up_w[expert][col] · x[token])`
///
/// Input contract identical to `dispatch_id_mm_for_test` except weight has
/// been split into `gate_w` and `up_w` (same shape, same expert stride).
///
/// Output `dst` is shape `[n_tokens * top_k, N]` (single fused result —
/// no intermediate `tmp_gate` / `tmp_up` buffers).
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
pub fn dispatch_id_mm_fused_gate_up_silu_for_test(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &MlxDevice,
    input: &MlxBuffer,
    gate_w: &MlxBuffer,
    up_w: &MlxBuffer,
    ids: &MlxBuffer,
    htpe: &mut MlxBuffer,
    hids: &mut MlxBuffer,
    output: &MlxBuffer,
    params: &GgmlIdMmDispatchParams,
) -> Result<()> {
    let (htpe, hids) = (&*htpe, &*hids);
    if input.dtype() != DType::F32
        || gate_w.dtype() != DType::U8
        || up_w.dtype() != DType::U8
        || ids.dtype() != DType::U32
        || htpe.dtype() != DType::U32
        || hids.dtype() != DType::U32
        || output.dtype() != DType::F32
    {
        return Err(MlxError::InvalidArgument(format!(
            "fused expert MM requires F32 input, native U8 gate/up GGUF blocks, U32 IDs/scratch, and F32 output; got {:?}/{:?}/{:?}/{:?}/{:?}/{:?}/{:?}",
            input.dtype(),
            gate_w.dtype(),
            up_w.dtype(),
            ids.dtype(),
            htpe.dtype(),
            hids.dtype(),
            output.dtype(),
        )));
    }
    // ---- Validate ----
    // Iter 4: Q6_K only. Q5_K + IQ4_XS variants ship in follow-up iters.
    match params.ggml_type {
        GgmlType::Q6_K => {}
        other => {
            return Err(MlxError::InvalidArgument(format!(
                "dispatch_id_mm_fused_gate_up_silu_for_test does not support {:?} (Q6_K only in iter 4)",
                other
            )));
        }
    }
    if params.n_tokens == 0
        || params.k == 0
        || params.n == 0
        || params.top_k == 0
        || params.n_experts == 0
    {
        return Err(MlxError::InvalidArgument(
            "n_tokens, K, N, top_k, n_experts must all be > 0".into(),
        ));
    }
    if params.top_k > params.n_experts {
        return Err(MlxError::InvalidArgument(
            "fused dispatch: top_k must not exceed n_experts".into(),
        ));
    }
    for (label, value) in [
        ("n_tokens", params.n_tokens),
        ("K", params.k),
        ("N", params.n),
        ("top_k", params.top_k),
        ("n_experts", params.n_experts),
    ] {
        if value > i32::MAX as u32 {
            return Err(MlxError::InvalidArgument(format!(
                "fused dispatch: {label} exceeds the signed Metal ABI"
            )));
        }
    }
    let qk = params.ggml_type.block_values();
    if params.k % qk != 0 {
        return Err(MlxError::InvalidArgument(format!(
            "K ({}) must be divisible by block QK ({})",
            params.k, qk
        )));
    }
    if params.top_k != 1 && params.top_k != 8 {
        return Err(MlxError::InvalidArgument(format!(
            "fused dispatch: no map0 instantiation for top_k={} (need 1 or 8)",
            params.top_k
        )));
    }

    let blocks_per_row = params.k / qk;
    let block_bytes = params.ggml_type.block_bytes();
    let expected_w = required_expert_weight_bytes(
        params.ggml_type,
        params.n_experts,
        params.n,
        params.k,
        params.expert_stride,
    )?;
    if gate_w.data_byte_len() < expected_w {
        return Err(MlxError::InvalidArgument(
            "fused dispatch: gate_w buffer too small".into(),
        ));
    }
    if up_w.data_byte_len() < expected_w {
        return Err(MlxError::InvalidArgument(
            "fused dispatch: up_w buffer too small".into(),
        ));
    }
    let input_bytes = checked_byte_extent(
        "fused mm input",
        &[
            params.n_tokens as usize,
            params.k as usize,
            DType::F32.size_of(),
        ],
    )?;
    if input.data_byte_len() < input_bytes {
        return Err(MlxError::InvalidArgument(
            "fused dispatch: input buffer too small".into(),
        ));
    }
    let total_rows = (params.n_tokens as usize)
        .checked_mul(params.top_k as usize)
        .ok_or_else(|| MlxError::InvalidArgument("fused mm expert row count overflow".into()))?;
    if total_rows > i32::MAX as usize {
        return Err(MlxError::InvalidArgument(
            "fused dispatch: n_tokens * top_k exceeds signed Metal row indexing".into(),
        ));
    }
    let ids_bytes = checked_byte_extent("fused mm ids", &[total_rows, DType::U32.size_of()])?;
    if ids.data_byte_len() < ids_bytes {
        return Err(MlxError::InvalidArgument(
            "fused dispatch: ids buffer too small".into(),
        ));
    }
    let output_bytes = checked_byte_extent(
        "fused mm output",
        &[total_rows, params.n as usize, DType::F32.size_of()],
    )?;
    if output.data_byte_len() < output_bytes {
        return Err(MlxError::InvalidArgument(
            "fused dispatch: output buffer too small".into(),
        ));
    }
    if htpe.data_byte_len() < params.htpe_bytes() {
        return Err(MlxError::InvalidArgument(
            "fused dispatch: htpe buffer too small".into(),
        ));
    }
    if hids.data_byte_len() < params.hids_bytes() {
        return Err(MlxError::InvalidArgument(
            "fused dispatch: hids buffer too small".into(),
        ));
    }

    // ---- Stage 1: map0 (identical to unfused path) ----
    let map0_kernel_name = match params.top_k {
        1 => "kernel_mul_mm_id_map0_ne20_1",
        8 => "kernel_mul_mm_id_map0_ne20_8",
        other => {
            return Err(MlxError::InvalidArgument(format!(
                "fused dispatch: no map0 instantiation for top_k={}",
                other
            )));
        }
    };
    let map0_pipeline = registry.get_pipeline(map0_kernel_name, device.metal_device())?;
    if u64::from(params.n_experts) > map0_pipeline.max_total_threads_per_threadgroup() {
        return Err(MlxError::InvalidArgument(format!(
            "fused dispatch: n_experts ({}) exceeds map0 pipeline threadgroup limit ({})",
            params.n_experts,
            map0_pipeline.max_total_threads_per_threadgroup(),
        )));
    }

    let map0_params = GgmlIdMmMap0GpuParams {
        ne10: params
            .n
            .try_into()
            .map_err(|_| MlxError::InvalidArgument("N out of i32 range".into()))?,
        ne11: params.top_k as i32,
        nb11: 0,
        nb12: 0,
        ne21: params.n_tokens as i32,
        ne20: params.top_k as i32,
        nb21: (params.top_k as u64) * (DType::U32.size_of() as u64),
    };

    let map0_shmem =
        (params.n_experts as u64) * (params.top_k as u64) * DType::U32.size_of() as u64;
    let map0_threadgroups = metal::MTLSize::new(1, 1, 1);
    let map0_threads = metal::MTLSize::new(params.n_experts as u64, 1, 1);

    encoder.encode_threadgroups_with_args_and_shared(
        map0_pipeline,
        &[
            (0, KernelArg::Bytes(as_bytes(&map0_params))),
            (1, KernelArg::Buffer(ids)),
            (2, KernelArg::Buffer(htpe)),
            (3, KernelArg::Buffer(hids)),
        ],
        &[(0, map0_shmem)],
        map0_threadgroups,
        map0_threads,
    );

    encoder.memory_barrier();

    // ---- Stage 2: fused mm_id ----
    let mm_pipeline = registry.get_pipeline(
        "kernel_fused_gate_up_silu_mm_id_q6_K_f32",
        device.metal_device(),
    )?;

    let nb01 = (blocks_per_row as u64) * (block_bytes as u64);
    let row_bytes = (params.k as u64) * (DType::F32.size_of() as u64);

    let mm_params = GgmlIdMmMmGpuParams {
        ne00: params.k as i32,
        ne02: params.n_experts as i32,
        nb01,
        nb02: params.expert_stride,
        nb03: 0,
        ne11: params.top_k as i32,
        _pad0: 0,
        nb10: DType::F32.size_of() as u64,
        nb11: 0,
        nb12: row_bytes,
        nb13: 0,
        ne20: params.top_k as i32,
        ne21: params.n_tokens as i32,
        ne0: params.n as i32,
        ne1: params.top_k as i32,
        r2: 1,
        r3: 1,
        sched_flags: 0,
        _pad1: 0,
    };

    const NR0: u64 = 64;
    const NR1: u64 = 32;
    const THREADS_PER_TG: u64 = 128;
    // Fused kernel doubles writeback buffer: 16 KB instead of 8 KB.
    const FUSED_SHMEM_BYTES: u64 = 16384;

    let mm_threadgroups = metal::MTLSize::new(
        (params.n_tokens as u64 + NR1 - 1) / NR1,
        (params.n as u64 + NR0 - 1) / NR0,
        params.n_experts as u64,
    );
    let mm_threads = metal::MTLSize::new(THREADS_PER_TG, 1, 1);

    encoder.encode_threadgroups_with_args_and_shared(
        mm_pipeline,
        &[
            (0, KernelArg::Bytes(as_bytes(&mm_params))),
            (1, KernelArg::Buffer(gate_w)),
            (2, KernelArg::Buffer(up_w)),
            (3, KernelArg::Buffer(input)),
            (4, KernelArg::Buffer(htpe)),
            (5, KernelArg::Buffer(hids)),
            (6, KernelArg::Buffer(output)),
        ],
        &[(0, FUSED_SHMEM_BYTES)],
        mm_threadgroups,
        mm_threads,
    );

    Ok(())
}

#[cfg(test)]
mod validation_tests {
    use super::*;

    #[test]
    fn padded_expert_stack_requires_last_stride_plus_matrix() {
        let matrix = ggml_expert_bytes(GgmlType::Q4_0, 1, 64, 256, 9_216).unwrap();
        assert_eq!(matrix, 9_216);
        assert_eq!(
            required_expert_weight_bytes(GgmlType::Q4_0, 3, 64, 256, 12_288).unwrap(),
            33_792
        );
    }

    #[test]
    fn expert_stack_rejects_stride_smaller_than_one_matrix() {
        let error = required_expert_weight_bytes(GgmlType::Q4_0, 3, 64, 256, 9_215)
            .expect_err("undersized stride must fail");
        assert!(error.to_string().contains("smaller than one matrix"));
    }

    #[test]
    fn mm_id_schedule_environment_values() {
        let default = GgmlRoutingPolicy::default().expert_mm_schedule;
        assert_eq!(default, ExpertMmSchedule::CompactTiles);
        for (value, expected, recognized) in [
            (None, default, true),
            (Some(""), default, true),
            (Some("  "), default, true),
            (Some("grid"), ExpertMmSchedule::UniformGrid, true),
            (Some("compact"), ExpertMmSchedule::CompactTiles, true),
            (Some(" grid "), ExpertMmSchedule::UniformGrid, true),
            (Some("GRID"), default, false),
            (Some("uniform"), default, false),
            (Some("1"), default, false),
        ] {
            assert_eq!(
                parse_mm_id_schedule(value),
                (expected, recognized),
                "HF2Q_MM_ID_SCHEDULE={value:?}"
            );
        }
    }
}