use half::bf16;
use std::collections::HashMap;
use std::ffi::c_void;
use xserv_kernels::*;
use xserv_tensor::{Device, Tensor};

use crate::config::ModelConfig;
use crate::paged_kv_cache::PagedKVCache;

pub struct GptOss {
    pub config: ModelConfig,
    embed_tokens: Tensor,
    layers: Vec<GptOssBlock>,
    norm: Tensor,
    norm_bias: Option<Tensor>,
    lm_head_t: Tensor,
    rope_cache: RopeCache,
    tp: Option<std::sync::Arc<xserv_distributed::TpContext>>,
    local_num_heads: usize,
    local_num_kv_heads: usize,
    has_norm_bias: bool,
}

struct GptOssBlock {
    input_norm: Tensor,
    input_norm_bias: Option<Tensor>,
    // Attention (with bias)
    q_proj_wt: Tensor,
    q_proj_bias: Tensor,
    k_proj_wt: Tensor,
    k_proj_bias: Tensor,
    v_proj_wt: Tensor,
    v_proj_bias: Tensor,
    o_proj_wt: Tensor,
    o_proj_bias: Tensor,
    sinks: Tensor,
    #[allow(dead_code)]
    is_sliding: bool,
    window_size: usize,
    // MoE MLP
    post_norm: Tensor,
    post_norm_bias: Option<Tensor>,
    router_wt: Tensor,
    router_bias: Tensor,
    // 3D expert weights for batched GEMM (contiguous on GPU)
    expert_gate_up_wt: Tensor,   // [local_experts, hidden, 2*inter] BF16
    expert_gate_up_bias: Tensor, // [local_experts, 2*inter]
    expert_down_wt: Tensor,      // [local_experts, inter, hidden] BF16
    expert_down_bias: Tensor,    // [local_experts, hidden]
    // FP8 quantized expert weights (Some when running FP8 W8A8)
    // Transposed layout [E, N, K] for cuBLASLt FP8 (Blackwell requires transA=T)
    expert_gate_up_fp8: Option<Tensor>, // [local_experts, 2*inter, hidden] FP8E4M3
    expert_gate_up_scale: Option<Tensor>, // [local_experts] F32
    expert_down_fp8: Option<Tensor>,    // [local_experts, hidden, inter] FP8E4M3
    expert_down_scale: Option<Tensor>,  // [local_experts] F32
    // MXFP4 W4A16 expert weights (Some when running 4-bit weight-only).
    // (packed [E, N, K/2] u8, scales [E, N, K/32] u8) in [E, N, K] layout.
    expert_gate_up_mxfp4: Option<(Tensor, Tensor)>,
    expert_down_mxfp4: Option<(Tensor, Tensor)>,
    local_experts: usize,
    // Activation params
    glu_alpha: f32,
    glu_limit: f32,
}

impl GptOss {
    pub fn from_weights(config: ModelConfig, w: HashMap<String, Tensor>) -> Self {
        Self::from_weights_tp(config, w, 0, 1, 0, None)
    }

    pub fn from_weights_tp(
        config: ModelConfig,
        mut w: HashMap<String, Tensor>,
        rank: usize,
        world: usize,
        device: u32,
        tp: Option<std::sync::Arc<xserv_distributed::TpContext>>,
    ) -> Self {
        crate::init_kernels();
        let dev = Device::Cuda(device);
        let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
            w.remove(name)
                .unwrap_or_else(|| panic!("missing weight: {name}"))
        };
        let repl = |t: Tensor| -> Tensor { t.to_device(dev) };
        // column-parallel: shard rows of [out, in], transpose → [in, out/world]
        let col = |t: Tensor| -> Tensor {
            shard_rows(&t, rank, world)
                .to_device(dev)
                .transpose(0, 1)
                .contiguous()
        };
        // row-parallel: shard cols of [out, in], transpose → [in/world, out]
        let row = |t: Tensor| -> Tensor {
            shard_cols(&t, rank, world)
                .to_device(dev)
                .transpose(0, 1)
                .contiguous()
        };
        // Bias sharding helpers
        let col_bias = |t: Tensor| -> Tensor { shard_1d(&t, rank, world).to_device(dev) };
        let repl_bias = |t: Tensor| -> Tensor { t.to_device(dev) };

        let embed_tokens = repl(take(&mut w, "model.embed_tokens.weight"));
        let norm = repl(take(&mut w, "model.norm.weight"));
        let norm_bias = w.remove("model.norm.bias").map(|t| repl(t));
        let lm_head_t = repl(take(&mut w, "lm_head.weight"))
            .transpose(0, 1)
            .contiguous();

        let head_dim = config.head_dim();
        let rope_theta = config.rope_theta.unwrap_or(150000.0);
        let max_seq_len = config.max_seq_len().min(8192); // cap for memory

        let rope_cache = if let Some(ref rs) = config.rope_scaling {
            if rs.rope_type.as_deref() == Some("yarn") {
                RopeCache::new_yarn(
                    max_seq_len,
                    head_dim,
                    rope_theta,
                    rs.factor.unwrap_or(1.0),
                    rs.original_max_position_embeddings.unwrap_or(4096),
                    rs.beta_fast.unwrap_or(32.0),
                    rs.beta_slow.unwrap_or(1.0),
                )
            } else {
                RopeCache::new(max_seq_len, head_dim, rope_theta as f32)
            }
        } else {
            RopeCache::new(max_seq_len, head_dim, rope_theta as f32)
        };

        let num_layers = config.num_layers();
        let num_experts = config.num_experts();
        let glu_alpha = config.geglu_alpha();
        let glu_limit = config.swiglu_limit.unwrap_or(7.0) as f32;

        let mut layers = Vec::with_capacity(num_layers);
        if rank == 0 {
            eprintln!(
                "Loading gpt-oss weights: {} layers, {} experts, world={world}, glu_alpha={glu_alpha}...",
                num_layers, num_experts
            );
        }

        for i in 0..num_layers {
            let p = format!("model.layers.{i}");

            // Attention weights — column-parallel for Q/K/V, row-parallel for O
            let q_proj_wt = col(take(&mut w, &format!("{p}.self_attn.q_proj.weight")));
            let q_proj_bias = col_bias(take(&mut w, &format!("{p}.self_attn.q_proj.bias")));
            let k_proj_wt = col(take(&mut w, &format!("{p}.self_attn.k_proj.weight")));
            let k_proj_bias = col_bias(take(&mut w, &format!("{p}.self_attn.k_proj.bias")));
            let v_proj_wt = col(take(&mut w, &format!("{p}.self_attn.v_proj.weight")));
            let v_proj_bias = col_bias(take(&mut w, &format!("{p}.self_attn.v_proj.bias")));
            let o_proj_wt = row(take(&mut w, &format!("{p}.self_attn.o_proj.weight")));
            let o_proj_bias = repl_bias(take(&mut w, &format!("{p}.self_attn.o_proj.bias")));

            // Sinks: shard per-head across TP ranks
            let sinks_full = take(&mut w, &format!("{p}.self_attn.sinks"));
            let sinks = shard_1d(&sinks_full, rank, world).to_device(dev);

            let is_sliding = config.is_sliding_layer(i);
            let window_size = if is_sliding { config.window_size() } else { 0 };

            // MoE weights — router replicated, experts split across TP ranks
            let router_wt_raw = take(&mut w, &format!("{p}.mlp.router.weight"));
            let router_wt = router_wt_raw.to_device(dev).transpose(0, 1).contiguous();
            let router_bias = repl_bias(take(&mut w, &format!("{p}.mlp.router.bias")));

            // Expert weights: [num_experts, hidden, 2*inter] — stored as 3D tensors
            // Expert parallelism: rank owns experts [rank*E/world .. (rank+1)*E/world)
            let gate_up_3d = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj"));
            let gate_up_bias_2d = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj_bias"));
            let down_3d = take(&mut w, &format!("{p}.mlp.experts.down_proj"));
            let down_bias_2d = take(&mut w, &format!("{p}.mlp.experts.down_proj_bias"));

            // FP8 scale tensors (present only in FP8-quantized models)
            let gate_up_scale = w.remove(&format!("{p}.mlp.experts.gate_up_proj_scale"));
            let down_scale = w.remove(&format!("{p}.mlp.experts.down_proj_scale"));

            let local_experts = num_experts / world;
            let expert_start = rank * local_experts;

            // MXFP4 stores 4-bit weights in an FP8E4M3 byte container (same dtype
            // as FP8), so distinguish by the scale rank: FP8 scale is 1-D [E],
            // MXFP4 scale is 3-D [E, N, K/32].
            let is_mxfp4 = gate_up_scale
                .as_ref()
                .map(|s| s.ndim() == 3)
                .unwrap_or(false);
            let is_fp8 = !is_mxfp4 && gate_up_3d.dtype() == xserv_tensor::DType::FP8E4M3;

            let mut expert_gate_up_mxfp4: Option<(Tensor, Tensor)> = None;
            let mut expert_down_mxfp4: Option<(Tensor, Tensor)> = None;

            let inter2 = if is_mxfp4 {
                gate_up_3d.shape()[1]
            } else {
                gate_up_3d.shape()[2]
            }; // 2*inter (N)
            let hidden = if is_mxfp4 {
                gate_up_3d.shape()[2] * 2
            } else {
                gate_up_3d.shape()[1]
            };
            let inter = if is_mxfp4 {
                down_3d.shape()[2] * 2
            } else {
                down_3d.shape()[1]
            };

            // Slice the rank's range of experts as contiguous 3D tensors on GPU
            let expert_gate_up_wt;
            let expert_down_wt;
            let expert_gate_up_fp8;
            let expert_gate_up_scale_gpu;
            let expert_down_fp8;
            let expert_down_scale_gpu;

            if is_mxfp4 {
                // MXFP4 W4A16: weights already [E, N, K] packed ([E, N, K/2] bytes)
                // + scales [E, N, K/32]. Slice this rank's experts (raw bytes).
                let gu_s = gate_up_scale.expect("MXFP4 model missing gate_up_proj_scale");
                let d_s = down_scale.expect("MXFP4 model missing down_proj_scale");
                let gu_packed = slice_expert_range_3d_raw(
                    &gate_up_3d,
                    expert_start,
                    local_experts,
                    inter2,
                    hidden / 2,
                )
                .to_device(dev);
                let gu_scl = slice_expert_range_3d_raw(
                    &gu_s,
                    expert_start,
                    local_experts,
                    inter2,
                    hidden / 32,
                )
                .to_device(dev);
                let dn_packed = slice_expert_range_3d_raw(
                    &down_3d,
                    expert_start,
                    local_experts,
                    hidden,
                    inter / 2,
                )
                .to_device(dev);
                let dn_scl = slice_expert_range_3d_raw(
                    &d_s,
                    expert_start,
                    local_experts,
                    hidden,
                    inter / 32,
                )
                .to_device(dev);
                expert_gate_up_mxfp4 = Some((gu_packed, gu_scl));
                expert_down_mxfp4 = Some((dn_packed, dn_scl));
                expert_gate_up_fp8 = None;
                expert_gate_up_scale_gpu = None;
                expert_down_fp8 = None;
                expert_down_scale_gpu = None;
                expert_gate_up_wt = Tensor::empty(&[1, 1, 1], xserv_tensor::DType::BF16, dev);
                expert_down_wt = Tensor::empty(&[1, 1, 1], xserv_tensor::DType::BF16, dev);
            } else if is_fp8 {
                // FP8 W8A8 path: load and TRANSPOSE weights for cuBLASLt (requires transA=T on Blackwell).
                // Original: [E, K, N] → Transposed: [E, N, K]
                let gu_sliced = slice_expert_range_3d_raw(
                    &gate_up_3d,
                    expert_start,
                    local_experts,
                    hidden,
                    inter2,
                );
                let dn_sliced =
                    slice_expert_range_3d_raw(&down_3d, expert_start, local_experts, inter, hidden);
                expert_gate_up_fp8 = Some(
                    transpose_3d_inner_raw(&gu_sliced, local_experts, hidden, inter2)
                        .to_device(dev),
                );
                expert_down_fp8 = Some(
                    transpose_3d_inner_raw(&dn_sliced, local_experts, inter, hidden).to_device(dev),
                );
                // Scales: [num_experts] F32 → slice to [local_experts]
                let gu_s = gate_up_scale.expect("FP8 model missing gate_up_proj_scale");
                let d_s = down_scale.expect("FP8 model missing down_proj_scale");
                expert_gate_up_scale_gpu =
                    Some(slice_scale_range(&gu_s, expert_start, local_experts).to_device(dev));
                expert_down_scale_gpu =
                    Some(slice_scale_range(&d_s, expert_start, local_experts).to_device(dev));
                // Dummy BF16 tensors (never read in FP8 path)
                expert_gate_up_wt = Tensor::empty(&[1, 1, 1], xserv_tensor::DType::BF16, dev);
                expert_down_wt = Tensor::empty(&[1, 1, 1], xserv_tensor::DType::BF16, dev);
            } else {
                // BF16 path: existing behavior
                expert_gate_up_wt =
                    slice_expert_range_3d(&gate_up_3d, expert_start, local_experts, hidden, inter2)
                        .to_device(dev);
                expert_down_wt =
                    slice_expert_range_3d(&down_3d, expert_start, local_experts, inter, hidden)
                        .to_device(dev);
                expert_gate_up_fp8 = None;
                expert_gate_up_scale_gpu = None;
                expert_down_fp8 = None;
                expert_down_scale_gpu = None;
            }
            let expert_gate_up_bias =
                slice_expert_range_2d(&gate_up_bias_2d, expert_start, local_experts, inter2)
                    .to_device(dev);
            let expert_down_bias =
                slice_expert_range_2d(&down_bias_2d, expert_start, local_experts, hidden)
                    .to_device(dev);

            xserv_cuda::allocator::cached_trim();

            let input_norm = repl(take(&mut w, &format!("{p}.input_layernorm.weight")));
            let input_norm_bias = w
                .remove(&format!("{p}.input_layernorm.bias"))
                .map(|t| repl(t));
            let post_norm = repl(take(
                &mut w,
                &format!("{p}.post_attention_layernorm.weight"),
            ));
            let post_norm_bias = w
                .remove(&format!("{p}.post_attention_layernorm.bias"))
                .map(|t| repl(t));

            layers.push(GptOssBlock {
                input_norm,
                input_norm_bias,
                q_proj_wt,
                q_proj_bias,
                k_proj_wt,
                k_proj_bias,
                v_proj_wt,
                v_proj_bias,
                o_proj_wt,
                o_proj_bias,
                sinks,
                is_sliding,
                window_size,
                post_norm,
                post_norm_bias,
                router_wt,
                router_bias,
                expert_gate_up_wt,
                expert_gate_up_bias,
                expert_down_wt,
                expert_down_bias,
                expert_gate_up_fp8,
                expert_gate_up_scale: expert_gate_up_scale_gpu,
                expert_down_fp8,
                expert_down_scale: expert_down_scale_gpu,
                expert_gate_up_mxfp4,
                expert_down_mxfp4,
                local_experts,
                glu_alpha,
                glu_limit,
            });
        }

        let local_num_heads = config.num_heads() / world;
        let local_num_kv_heads = config.num_kv_heads() / world;

        let has_norm_bias = norm_bias.is_some();
        let is_fp8 = layers
            .first()
            .map(|l| l.expert_gate_up_fp8.is_some())
            .unwrap_or(false);
        let is_mxfp4 = layers
            .first()
            .map(|l| l.expert_gate_up_mxfp4.is_some())
            .unwrap_or(false);
        if rank == 0 {
            if has_norm_bias {
                eprintln!("gpt-oss: detected LayerNorm bias — using LayerNorm instead of RMSNorm");
            }
            if is_fp8 {
                eprintln!(
                    "gpt-oss: FP8 E4M3 quantized expert weights detected (W8A8 cuBLASLt mode)"
                );
            }
            if is_mxfp4 {
                eprintln!(
                    "gpt-oss: MXFP4 quantized expert weights detected (W4A16 fused-GEMV mode)"
                );
            }
        }

        // Warn about unused weights that the model didn't consume
        if rank == 0 && !w.is_empty() {
            eprintln!("WARNING: {} unused weight(s) in model:", w.len());
            let mut keys: Vec<_> = w.keys().collect();
            keys.sort();
            for k in &keys {
                eprintln!("  {k}");
            }
        }

        Self {
            config,
            embed_tokens,
            layers,
            norm,
            norm_bias,
            lm_head_t,
            rope_cache,
            tp,
            local_num_heads,
            local_num_kv_heads,
            has_norm_bias,
        }
    }

    #[inline]
    fn all_reduce(&self, t: &Tensor) {
        if let Some(tp) = &self.tp {
            if tp.world > 1 {
                let ptr = t.storage().gpu_buffer().as_ptr() as *mut c_void;
                tp.all_reduce_sum_bf16_ptr(ptr, t.numel());
            }
        }
    }

    #[inline]
    fn norm(x: &Tensor, weight: &Tensor, bias: &Option<Tensor>, eps: f32) -> Tensor {
        match bias {
            Some(b) => layernorm(x, weight, b, eps),
            None => rmsnorm(x, weight, eps),
        }
    }

    #[inline]
    fn add_norm(
        x: &Tensor,
        residual: &Tensor,
        weight: &Tensor,
        bias: &Option<Tensor>,
        eps: f32,
    ) -> (Tensor, Tensor) {
        match bias {
            Some(b) => {
                let sum = xserv_kernels::add(x, residual);
                let normed = layernorm(&sum, weight, b, eps);
                (normed, sum)
            }
            None => add_rmsnorm(x, residual, weight, eps),
        }
    }

    fn norm_eps(&self) -> f32 {
        if self.has_norm_bias {
            self.config.ln_eps()
        } else {
            self.config.rms_norm_eps.unwrap_or(1e-5) as f32
        }
    }

    /// Paged decode: process one token per sequence using paged KV cache.
    pub fn forward_decode_paged(
        &self,
        tokens: &[u32],
        positions: &[usize],
        seq_slots: &[usize],
        paged_cache: &mut PagedKVCache,
    ) -> Tensor {
        let batch = tokens.len();
        assert_eq!(positions.len(), batch);
        assert_eq!(seq_slots.len(), batch);
        assert!(batch > 0);

        self.decode_prepare(positions, seq_slots, paged_cache);

        // Upload token ids + positions, then run the pure-GPU core.
        let ids_gpu = upload_u32(tokens);
        let positions_u32: Vec<u32> = positions.iter().map(|&p| p as u32).collect();
        let pos_gpu = upload_u32(&positions_u32);
        let logits = self.decode_core(
            ids_gpu.as_ptr() as *const c_void,
            pos_gpu.as_ptr() as *const c_void,
            batch,
            paged_cache,
        );

        for &slot in seq_slots {
            paged_cache.advance_seq_len(slot, 1);
        }
        logits
    }

    /// Host-side per-step cache bookkeeping: block allocation + uploading
    /// block tables / context lens to their (stable-address) GPU buffers.
    /// Runs OUTSIDE the CUDA-graph captured region.
    pub fn decode_prepare(
        &self,
        positions: &[usize],
        seq_slots: &[usize],
        paged_cache: &mut PagedKVCache,
    ) {
        let kv_lens: Vec<i32> = positions.iter().map(|&p| (p + 1) as i32).collect();
        for (b, &slot) in seq_slots.iter().enumerate() {
            paged_cache.ensure_capacity(slot, positions[b] + 1);
        }
        paged_cache.sync_active_batch_with_lens(seq_slots, &kv_lens);
    }

    /// The pure-GPU decode step: embedding → 24 layers → final norm → logits.
    /// Token ids and positions are read from device buffers; every other input
    /// (weights, KV pools, block table, context lens) has a stable address —
    /// which is exactly what makes this region CUDA-graph capturable.
    pub fn decode_core(
        &self,
        ids_gpu: *const c_void,
        pos_gpu: *const c_void,
        batch: usize,
        paged_cache: &mut PagedKVCache,
    ) -> Tensor {
        let num_heads = self.local_num_heads;
        let num_kv_heads = self.local_num_kv_heads;
        let head_dim = self.config.head_dim();
        let eps = self.norm_eps();

        let bt_ptr = paged_cache.block_table_gpu().as_ptr() as *const i32;
        let cl_ptr = paged_cache.context_lens_gpu().as_ptr() as *const i32;
        let max_blocks = paged_cache.max_blocks_per_seq();

        let mut x = embedding_device_ids(&self.embed_tokens, ids_gpu, batch);

        for (layer_idx, layer) in self.layers.iter().enumerate() {
            let residual = x.clone();
            let normed = Self::norm(&x, &layer.input_norm, &layer.input_norm_bias, eps);

            // Q/K/V projections with bias
            let q_all = add_bias(&matmul_2d(&normed, &layer.q_proj_wt), &layer.q_proj_bias);
            let k_all = add_bias(&matmul_2d(&normed, &layer.k_proj_wt), &layer.k_proj_bias);
            let v_all = add_bias(&matmul_2d(&normed, &layer.v_proj_wt), &layer.v_proj_bias);

            // Reshape for RoPE: [B, H*D] → [B, H, D]
            let q_3d = q_all.reshape(&[batch, num_heads, head_dim]);
            let k_3d = k_all.reshape(&[batch, num_kv_heads, head_dim]);

            // RoPE (no QK-norm for gpt-oss)
            rope_inplace_device_pos(&q_3d, &self.rope_cache, pos_gpu);
            rope_inplace_device_pos(&k_3d, &self.rope_cache, pos_gpu);

            let v_3d = v_all.reshape(&[batch, num_kv_heads, head_dim]);

            // KV cache scatter
            paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, batch);

            // Paged attention with sinks + sliding window
            let q_4d = q_3d.reshape(&[batch, num_heads, 1, head_dim]);
            let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const c_void;
            let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const c_void;
            let sinks_ptr = layer.sinks.data_ptr() as *const c_void;

            let attn_out = paged_decode_attention_sinks(
                &q_4d,
                k_pool_ptr,
                v_pool_ptr,
                bt_ptr,
                cl_ptr,
                sinks_ptr,
                batch,
                num_heads,
                num_kv_heads,
                head_dim,
                max_blocks,
                layer.window_size,
            );

            let attn_merged = attn_out.reshape(&[batch, num_heads * head_dim]);
            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
            self.all_reduce(&attn_proj);
            let attn_proj = add_bias(&attn_proj, &layer.o_proj_bias);

            // Residual + post-norm
            let (normed, x_new) = Self::add_norm(
                &attn_proj,
                &residual,
                &layer.post_norm,
                &layer.post_norm_bias,
                eps,
            );

            let residual = x_new;
            let normed = normed.contiguous();

            // MoE MLP
            let moe_out = self.moe_forward(&normed, layer, batch);
            x = xserv_kernels::add(&residual, &moe_out);
        }

        let x = Self::norm(&x, &self.norm, &self.norm_bias, eps);
        matmul_2d(&x, &self.lm_head_t)
    }

    /// Paged prefill: process full prompt tokens.
    pub fn forward_prefill_paged(
        &self,
        token_ids: &[u32],
        slot: usize,
        paged_cache: &mut PagedKVCache,
    ) -> Tensor {
        let new_tokens = token_ids.len();
        let pos_offset = paged_cache.seq_len(slot);
        let num_heads = self.local_num_heads;
        let num_kv_heads = self.local_num_kv_heads;
        let head_dim = self.config.head_dim();
        let eps = self.norm_eps();

        paged_cache.ensure_capacity(slot, pos_offset + new_tokens);
        paged_cache.advance_seq_len(slot, new_tokens);

        let mut x = embedding(&self.embed_tokens, token_ids);
        let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens)
            .map(|p| p as u32)
            .collect();

        for (layer_idx, layer) in self.layers.iter().enumerate() {
            let residual = x.clone();
            let normed = Self::norm(&x, &layer.input_norm, &layer.input_norm_bias, eps);

            let q = add_bias(&matmul_2d(&normed, &layer.q_proj_wt), &layer.q_proj_bias);
            let k = add_bias(&matmul_2d(&normed, &layer.k_proj_wt), &layer.k_proj_bias);
            let v = add_bias(&matmul_2d(&normed, &layer.v_proj_wt), &layer.v_proj_bias);

            let q = reshape_heads_gpu(&q, new_tokens, num_heads, head_dim);
            let k = reshape_heads_gpu(&k, new_tokens, num_kv_heads, head_dim);
            let v = reshape_heads_gpu(&v, new_tokens, num_kv_heads, head_dim);

            // RoPE
            let q = transpose_for_rope_gpu(&q, new_tokens, num_heads, head_dim);
            let k = transpose_for_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
            rope_inplace(&q, &self.rope_cache, &positions);
            rope_inplace(&k, &self.rope_cache, &positions);
            let q = transpose_from_rope_gpu(&q, new_tokens, num_heads, head_dim);
            let k = transpose_from_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);

            // KV cache
            paged_cache.append_tokens(slot, layer_idx, &k, &v, new_tokens, pos_offset);
            let (k_full, v_full) = paged_cache.gather_kv_contiguous(slot, layer_idx);

            // Flash attention with gpt-oss sinks + (per-layer) sliding window.
            let attn_out =
                flash_attention_sinks(&q, &k_full, &v_full, &layer.sinks, layer.window_size);

            let attn_merged = merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
            self.all_reduce(&attn_proj);
            let attn_proj = add_bias(&attn_proj, &layer.o_proj_bias);

            let (normed, x_new) = Self::add_norm(
                &attn_proj,
                &residual,
                &layer.post_norm,
                &layer.post_norm_bias,
                eps,
            );
            let residual = x_new;

            // MoE MLP
            let moe_out = self.moe_forward(&normed, layer, new_tokens);
            x = xserv_kernels::add(&residual, &moe_out);
        }

        let x = Self::norm(&x, &self.norm, &self.norm_bias, eps);
        matmul_2d(&x, &self.lm_head_t)
    }

    /// MoE forward pass — fully on GPU via batched GEMM.
    ///
    /// Each rank owns `local_experts` experts. The input is replicated across all
    /// local experts, processed via two batched cuBLAS GEMMs (gate_up and down),
    /// and the selected experts' outputs are weighted-summed on GPU. Non-selected
    /// experts contribute zero (via the routing weights), so no scatter/gather is
    /// needed. AllReduce sums partial results across TP ranks.
    fn moe_forward(&self, x: &Tensor, layer: &GptOssBlock, num_tokens: usize) -> Tensor {
        let num_experts = self.config.num_experts();
        let top_k = self.config.experts_per_token();
        let rank = self.tp.as_ref().map(|tp| tp.rank).unwrap_or(0);
        let local_experts = layer.local_experts;
        let expert_start = rank * local_experts;

        // 1. Router: [tokens, hidden] @ [hidden, num_experts] + bias → [tokens, num_experts]
        let router_logits = add_bias(&matmul_2d(x, &layer.router_wt), &layer.router_bias);

        // 2. GPU top-k + softmax
        let (topk_ids, topk_weights) =
            xserv_kernels::moe::moe_topk_softmax(&router_logits, num_experts, top_k);

        // Sparse decode path: compute ONLY the routed experts. The dense path
        // below reads every local expert's weights per forward; the sparse
        // GEMVs read ~top_k/num_experts of the bytes, which dominates decode
        // (memory-bound). Dense reads each weight once for ALL tokens, so it
        // wins back at num_tokens ≈ local_experts / E[local hits] ≈ 8.
        const SPARSE_MAX_TOKENS: usize = 8;
        let quantized = layer.expert_gate_up_fp8.is_some() || layer.expert_gate_up_mxfp4.is_some();
        if num_tokens <= SPARSE_MAX_TOKENS && quantized && !dense_moe_forced() {
            let gate_up = if let Some((ref packed, ref scales)) = layer.expert_gate_up_mxfp4 {
                let n = packed.shape()[1];
                let k = packed.shape()[2] * 2;
                xserv_kernels::moe::moe_sparse_gemv_mxfp4(
                    x,
                    packed,
                    scales,
                    &layer.expert_gate_up_bias,
                    &topk_ids,
                    num_tokens,
                    top_k,
                    n,
                    k,
                    expert_start,
                    local_experts,
                    false,
                )
            } else {
                xserv_kernels::moe::moe_sparse_gemv_fp8(
                    x,
                    layer.expert_gate_up_fp8.as_ref().unwrap(),
                    layer.expert_gate_up_scale.as_ref().unwrap(),
                    &layer.expert_gate_up_bias,
                    &topk_ids,
                    num_tokens,
                    top_k,
                    expert_start,
                    local_experts,
                    false,
                )
            };

            // GLU over all slots. Non-local slots hold unwritten memory; they
            // are never consumed (the down GEMV and the weighted sum both skip
            // slots whose expert this rank does not own).
            let inter2 = gate_up.shape()[2];
            let gate_up_flat = gate_up.reshape(&[num_tokens * top_k, inter2]);
            let activated = gpt_oss_glu(&gate_up_flat, layer.glu_alpha, layer.glu_limit);

            let down = if let Some((ref packed, ref scales)) = layer.expert_down_mxfp4 {
                let n = packed.shape()[1];
                let k = packed.shape()[2] * 2;
                xserv_kernels::moe::moe_sparse_gemv_mxfp4(
                    &activated,
                    packed,
                    scales,
                    &layer.expert_down_bias,
                    &topk_ids,
                    num_tokens,
                    top_k,
                    n,
                    k,
                    expert_start,
                    local_experts,
                    true,
                )
            } else {
                xserv_kernels::moe::moe_sparse_gemv_fp8(
                    &activated,
                    layer.expert_down_fp8.as_ref().unwrap(),
                    layer.expert_down_scale.as_ref().unwrap(),
                    &layer.expert_down_bias,
                    &topk_ids,
                    num_tokens,
                    top_k,
                    expert_start,
                    local_experts,
                    true,
                )
            };

            let moe_out = xserv_kernels::moe::moe_weighted_sum_sparse(
                &down,
                &topk_ids,
                &topk_weights,
                expert_start,
                local_experts,
            );
            self.all_reduce(&moe_out);
            return moe_out;
        }

        // 3. Replicate input: [tokens, hidden] → [local_experts, tokens, hidden]
        let x_rep = xserv_kernels::moe::moe_replicate(x, local_experts);

        // 4. Batched GEMM gate_up: [E, tokens, hidden] @ [E, hidden, 2*inter] → [E, tokens, 2*inter]
        let gate_up = if let Some((ref packed, ref scales)) = layer.expert_gate_up_mxfp4 {
            // MXFP4 W4A16: decode (M=1) uses the fused 4-bit dequant GEMV; prefill
            // dequantizes to BF16 then reuses the batched GEMM.
            let n = packed.shape()[1];
            let k = packed.shape()[2] * 2;
            if num_tokens == 1 {
                let x2 = x_rep.reshape(&[local_experts, k]);
                xserv_kernels::quantization::batched_gemv_mxfp4(&x2, packed, scales, n, k)
                    .reshape(&[local_experts, 1, n])
            } else {
                let w_bf16 = xserv_kernels::quantization::dequant_mxfp4_to_bf16_t(
                    packed,
                    scales,
                    local_experts,
                    n,
                    k,
                );
                xserv_kernels::moe::batched_gemm_strided(&x_rep, &w_bf16)
            }
        } else if let Some(ref wt_fp8_t) = layer.expert_gate_up_fp8 {
            // W8A8: quantize activations with per-expert scalar scale, use cuBLASLt FP8 GEMM
            let (x_fp8, x_scales) =
                xserv_kernels::quantization::quantize_bf16_to_fp8_rowwise(&x_rep);
            xserv_kernels::quantization::batched_gemm_fp8(
                &x_fp8,
                &x_scales,
                wt_fp8_t,
                layer.expert_gate_up_scale.as_ref().unwrap(),
            )
        } else {
            xserv_kernels::moe::batched_gemm_strided(&x_rep, &layer.expert_gate_up_wt)
        };

        // 5. Bias add: gate_up += expert_gate_up_bias  (in-place)
        xserv_kernels::moe::moe_bias_add_3d(&gate_up, &layer.expert_gate_up_bias);

        // 6. GLU activation: treat [E * tokens, 2*inter] → [E * tokens, inter]
        let inter2 = gate_up.shape()[2];
        let flat_rows = local_experts * num_tokens;
        let gate_up_flat = gate_up.reshape(&[flat_rows, inter2]);
        let activated_flat = gpt_oss_glu(&gate_up_flat, layer.glu_alpha, layer.glu_limit);
        let inter = inter2 / 2;
        let activated = activated_flat.reshape(&[local_experts, num_tokens, inter]);

        // 7. Batched GEMM down: [E, tokens, inter] @ [E, inter, hidden] → [E, tokens, hidden]
        let down = if let Some((ref packed, ref scales)) = layer.expert_down_mxfp4 {
            let n = packed.shape()[1];
            let k = packed.shape()[2] * 2;
            if num_tokens == 1 {
                let a2 = activated.reshape(&[local_experts, k]);
                xserv_kernels::quantization::batched_gemv_mxfp4(&a2, packed, scales, n, k)
                    .reshape(&[local_experts, 1, n])
            } else {
                let w_bf16 = xserv_kernels::quantization::dequant_mxfp4_to_bf16_t(
                    packed,
                    scales,
                    local_experts,
                    n,
                    k,
                );
                xserv_kernels::moe::batched_gemm_strided(&activated, &w_bf16)
            }
        } else if let Some(ref wt_fp8) = layer.expert_down_fp8 {
            // W8A8: quantize post-GLU activations to FP8, use cuBLASLt FP8 GEMM
            let (act_fp8, act_scales) =
                xserv_kernels::quantization::quantize_bf16_to_fp8_rowwise(&activated);
            xserv_kernels::quantization::batched_gemm_fp8(
                &act_fp8,
                &act_scales,
                wt_fp8,
                layer.expert_down_scale.as_ref().unwrap(),
            )
        } else {
            xserv_kernels::moe::batched_gemm_strided(&activated, &layer.expert_down_wt)
        };

        // 8. Bias add: down += expert_down_bias  (in-place)
        xserv_kernels::moe::moe_bias_add_3d(&down, &layer.expert_down_bias);

        // 9. Weighted sum across experts → [tokens, hidden]
        let moe_out = xserv_kernels::moe::moe_weighted_sum(
            &down,
            &topk_ids,
            &topk_weights,
            expert_start,
            local_experts,
            top_k,
        );

        self.all_reduce(&moe_out);
        moe_out
    }
}

// --- Helpers ---

/// Upload a u32 slice to a pooled GPU buffer (synchronous H2D).
fn upload_u32(vals: &[u32]) -> xserv_cuda::GpuBuffer {
    let bytes = unsafe { std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4) };
    let mut buf = xserv_cuda::allocator::cached_alloc(bytes.len()).expect("alloc u32 upload");
    buf.copy_from_host(bytes).unwrap();
    buf
}

/// XSERV_DENSE_MOE=1 forces the dense all-expert path (A/B benchmarking).
fn dense_moe_forced() -> bool {
    static FORCED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *FORCED.get_or_init(|| std::env::var("XSERV_DENSE_MOE").is_ok_and(|v| v != "0"))
}

fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
    assert_eq!(a.ndim(), 2);
    assert_eq!(b.ndim(), 2);
    matmul(a, b, GemmBackend::CuBlas)
}

/// Add bias to a 2D tensor: [rows, cols] + [cols] → [rows, cols].
/// Single GPU broadcast kernel — the old rows>1 path tiled the bias on the
/// CPU (D2H + host loop + H2D) on every call, 96×/prefill in the hot path.
fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
    let x_c = x.contiguous();
    xserv_kernels::bias_add_2d(&x_c, bias)
}

fn shard_rows(t: &Tensor, rank: usize, world: usize) -> Tensor {
    if world == 1 {
        return t.clone();
    }
    let shape = t.shape();
    assert_eq!(shape.len(), 2);
    let (rows, cols) = (shape[0], shape[1]);
    assert!(
        rows % world == 0,
        "rows {rows} not divisible by world {world}"
    );
    let local = rows / world;
    let host = t.to_device(Device::Cpu);
    let data = host.as_slice::<bf16>();
    let start = rank * local * cols;
    let shard = data[start..start + local * cols].to_vec();
    Tensor::from_slice(&shard, &[local, cols])
}

fn shard_cols(t: &Tensor, rank: usize, world: usize) -> Tensor {
    if world == 1 {
        return t.clone();
    }
    let shape = t.shape();
    assert_eq!(shape.len(), 2);
    let (rows, cols) = (shape[0], shape[1]);
    assert!(
        cols % world == 0,
        "cols {cols} not divisible by world {world}"
    );
    let local = cols / world;
    let c0 = rank * local;
    let host = t.to_device(Device::Cpu);
    let data = host.as_slice::<bf16>();
    let mut shard = Vec::with_capacity(rows * local);
    for r in 0..rows {
        let base = r * cols + c0;
        shard.extend_from_slice(&data[base..base + local]);
    }
    Tensor::from_slice(&shard, &[rows, local])
}

fn shard_1d(t: &Tensor, rank: usize, world: usize) -> Tensor {
    if world == 1 {
        return t.clone();
    }
    let shape = t.shape();
    assert_eq!(shape.len(), 1);
    let total = shape[0];
    assert!(
        total % world == 0,
        "dim {total} not divisible by world {world}"
    );
    let local = total / world;
    let host = t.to_device(Device::Cpu);
    let data = host.as_slice::<bf16>();
    let start = rank * local;
    let shard = data[start..start + local].to_vec();
    Tensor::from_slice(&shard, &[local])
}

/// Transpose the inner two dimensions of a [batch, rows, cols] tensor → [batch, cols, rows].
/// Works on raw bytes (any dtype). CPU-only.
fn transpose_3d_inner_raw(t: &Tensor, batch: usize, rows: usize, cols: usize) -> Tensor {
    assert_eq!(t.ndim(), 3);
    assert_eq!(t.shape(), &[batch, rows, cols]);
    let host = t.to_device(Device::Cpu);
    let es = t.dtype().size_bytes();
    let raw = host.as_raw_bytes();
    let mut out = vec![0u8; batch * cols * rows * es];
    for b in 0..batch {
        for r in 0..rows {
            for c in 0..cols {
                let src_off = (b * rows * cols + r * cols + c) * es;
                let dst_off = (b * cols * rows + c * rows + r) * es;
                out[dst_off..dst_off + es].copy_from_slice(&raw[src_off..src_off + es]);
            }
        }
    }
    Tensor::from_raw_bytes(&out, &[batch, cols, rows], t.dtype())
}

/// Extract experts [start..start+count) from a [num_experts, rows, cols] 3D tensor (any dtype, raw bytes).
fn slice_expert_range_3d_raw(
    t: &Tensor,
    start: usize,
    count: usize,
    rows: usize,
    cols: usize,
) -> Tensor {
    assert_eq!(t.ndim(), 3);
    let host = t.to_device(Device::Cpu);
    let elem_size = t.dtype().size_bytes();
    let raw = host.as_raw_bytes();
    let stride = rows * cols * elem_size;
    let offset = start * stride;
    let slice = &raw[offset..offset + count * stride];
    Tensor::from_raw_bytes(slice, &[count, rows, cols], t.dtype())
}

/// Slice scale tensor [num_experts] F32 → [count] starting at `start`.
fn slice_scale_range(t: &Tensor, start: usize, count: usize) -> Tensor {
    assert_eq!(t.ndim(), 1);
    assert_eq!(t.dtype(), xserv_tensor::DType::F32);
    let host = t.to_device(Device::Cpu);
    let data = host.as_slice::<f32>();
    let slice = data[start..start + count].to_vec();
    Tensor::from_slice(&slice, &[count])
}

/// Extract experts [start..start+count) from a [num_experts, rows, cols] 3D tensor
fn slice_expert_range_3d(
    t: &Tensor,
    start: usize,
    count: usize,
    rows: usize,
    cols: usize,
) -> Tensor {
    assert_eq!(t.ndim(), 3);
    let host = t.to_device(Device::Cpu);
    let data = host.as_slice::<bf16>();
    let stride = rows * cols;
    let offset = start * stride;
    let slice = data[offset..offset + count * stride].to_vec();
    Tensor::from_slice(&slice, &[count, rows, cols])
}

/// Extract experts [start..start+count) from a [num_experts, dim] 2D tensor
fn slice_expert_range_2d(t: &Tensor, start: usize, count: usize, dim: usize) -> Tensor {
    assert_eq!(t.ndim(), 2);
    let host = t.to_device(Device::Cpu);
    let data = host.as_slice::<bf16>();
    let offset = start * dim;
    let slice = data[offset..offset + count * dim].to_vec();
    Tensor::from_slice(&slice, &[count, dim])
}