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

use crate::config::ModelConfig;
use crate::gpt2::KVCache;
use crate::kv_cache::GpuKVCache;
use crate::paged_kv_cache::PagedKVCache;

pub struct Qwen3 {
    pub config: ModelConfig,
    embed_tokens: Tensor,
    layers: Vec<Qwen3Block>,
    norm: Tensor,
    lm_head_t: Tensor, // precomputed transpose
    rope_cache: RopeCache,
    // Tensor parallelism. `tp` is None (or world==1) for single-GPU; otherwise
    // this rank holds 1/world of the heads and AllReduces after o_proj/down_proj.
    tp: Option<std::sync::Arc<xserv_distributed::TpContext>>,
    local_num_heads: usize,    // = num_heads / world
    local_num_kv_heads: usize, // = num_kv_heads / world
    // Pipeline parallelism (Phase 18): this stage holds a contiguous slice of
    // layers. `is_first_stage` owns `embed_tokens`; `is_last_stage` owns
    // `norm`/`lm_head_t`. Both true for single-GPU / TP (the whole model).
    is_first_stage: bool,
    is_last_stage: bool,
}

struct Qwen3Block {
    input_norm: Tensor,      // [hidden]
    qkv_proj_wt: Tensor,     // FUSED: [hidden, (H+2*KV)*D] — Q|K|V columns
    q_dim: usize,            // num_heads * head_dim (Q slice boundary)
    kv_dim: usize,           // num_kv_heads * head_dim (K/V slice size)
    o_proj_wt: Tensor,       // TRANSPOSED: [num_heads*head_dim, hidden]
    q_norm: Tensor,          // [head_dim]
    k_norm: Tensor,          // [head_dim]
    post_norm: Tensor,       // [hidden]
    gate_up_proj_wt: Tensor, // FUSED: [hidden, 2*intermediate]
    down_proj_wt: Tensor,    // TRANSPOSED: [intermediate, hidden]
}

impl Qwen3Block {
    fn q_proj_wt(&self) -> Tensor {
        self.qkv_proj_wt.narrow(1, 0, self.q_dim)
    }
    fn k_proj_wt(&self) -> Tensor {
        self.qkv_proj_wt.narrow(1, self.q_dim, self.kv_dim)
    }
    fn v_proj_wt(&self) -> Tensor {
        self.qkv_proj_wt
            .narrow(1, self.q_dim + self.kv_dim, self.kv_dim)
    }
    fn gate_proj_wt(&self) -> Tensor {
        let half = self.gate_up_proj_wt.shape()[1] / 2;
        self.gate_up_proj_wt.narrow(1, 0, half)
    }
    fn up_proj_wt(&self) -> Tensor {
        let half = self.gate_up_proj_wt.shape()[1] / 2;
        self.gate_up_proj_wt.narrow(1, half, half)
    }
}

impl Qwen3 {
    /// Single-GPU load (weights already on the target GPU). Equivalent to
    /// `from_weights_tp(.., rank=0, world=1, device=0, tp=None)`.
    pub fn from_weights(config: ModelConfig, w: HashMap<String, Tensor>) -> Self {
        Self::from_weights_tp(config, w, 0, 1, 0, None)
    }

    /// Tensor-parallel load. `w` may live on CPU or any device; each weight is
    /// sharded for `rank`/`world`, uploaded to `device`, and transposed.
    /// `world==1` shards are identity, so this is also the single-GPU path.
    ///
    /// Split scheme (Megatron-style):
    /// - column-parallel (split output): q/k/v/gate/up  → shard rows of `[out,in]`
    /// - row-parallel    (split input):  o/down         → shard cols of `[out,in]`
    /// - replicated: norms, embed_tokens, lm_head
    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}"))
        };
        // Replicated weight: upload whole to this rank's device.
        let repl = |t: Tensor| -> Tensor { t.to_device(dev) };
        // column-parallel: keep this rank's rows of [out, in], upload, transpose → [in, out/world].
        let col = |t: Tensor| -> Tensor {
            shard_rows(&t, rank, world)
                .to_device(dev)
                .transpose(0, 1)
                .contiguous()
        };
        // row-parallel: keep this rank's cols of [out, in], upload, transpose → [in/world, out].
        let row = |t: Tensor| -> Tensor {
            shard_cols(&t, rank, world)
                .to_device(dev)
                .transpose(0, 1)
                .contiguous()
        };

        let embed_tokens = repl(take(&mut w, "model.embed_tokens.weight"));
        let norm = repl(take(&mut w, "model.norm.weight"));
        let lm_head_t = repl(take(&mut w, "lm_head.weight"))
            .transpose(0, 1)
            .contiguous();

        let rope_cache = RopeCache::new(
            config.max_seq_len(),
            config.head_dim(),
            config.rope_theta.unwrap_or(1_000_000.0) as f32,
        );

        let num_layers = config.num_layers();
        let mut layers = Vec::with_capacity(num_layers);
        if rank == 0 {
            eprintln!(
                "Loading+sharding weights for {} layers (world={world})...",
                num_layers
            );
        }
        for i in 0..num_layers {
            let p = format!("model.layers.{i}");
            let q_proj_wt = col(take(&mut w, &format!("{p}.self_attn.q_proj.weight")));
            let k_proj_wt = col(take(&mut w, &format!("{p}.self_attn.k_proj.weight")));
            let v_proj_wt = col(take(&mut w, &format!("{p}.self_attn.v_proj.weight")));
            let q_dim = q_proj_wt.shape()[1];
            let kv_dim = k_proj_wt.shape()[1];
            let qkv_proj_wt = cat_cols(&[&q_proj_wt, &k_proj_wt, &v_proj_wt]);
            drop((q_proj_wt, k_proj_wt, v_proj_wt));
            let gate_proj_wt = col(take(&mut w, &format!("{p}.mlp.gate_proj.weight")));
            let up_proj_wt = col(take(&mut w, &format!("{p}.mlp.up_proj.weight")));
            let gate_up_proj_wt = cat_cols(&[&gate_proj_wt, &up_proj_wt]);
            drop((gate_proj_wt, up_proj_wt));
            xserv_cuda::allocator::cached_trim();
            layers.push(Qwen3Block {
                input_norm: repl(take(&mut w, &format!("{p}.input_layernorm.weight"))),
                qkv_proj_wt,
                q_dim,
                kv_dim,
                o_proj_wt: row(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
                q_norm: repl(take(&mut w, &format!("{p}.self_attn.q_norm.weight"))),
                k_norm: repl(take(&mut w, &format!("{p}.self_attn.k_norm.weight"))),
                post_norm: repl(take(
                    &mut w,
                    &format!("{p}.post_attention_layernorm.weight"),
                )),
                gate_up_proj_wt,
                down_proj_wt: row(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
            });
        }

        Self {
            local_num_heads: config.num_heads() / world,
            local_num_kv_heads: config.num_kv_heads() / world,
            config,
            embed_tokens,
            layers,
            norm,
            lm_head_t,
            rope_cache,
            tp,
            is_first_stage: true,
            is_last_stage: true,
        }
    }

    /// Pipeline-parallel load (Phase 18). This stage holds the contiguous layer
    /// range `[stage*L, (stage+1)*L)` with `L = num_layers / num_stages`; only
    /// stage 0 keeps `embed_tokens` and only the last stage keeps `norm`/`lm_head`
    /// (others get a 1x1 placeholder, guarded by the stage flags and never used).
    /// Heads are NOT split (PP is orthogonal to TP), so each stage runs full
    /// attention/MLP over its layers and hands off the `[tokens, hidden]` hidden
    /// state to the next stage (the engine does the NCCL send/recv).
    pub fn from_weights_pp(
        config: ModelConfig,
        mut w: HashMap<String, Tensor>,
        stage: usize,
        num_stages: usize,
        device: u32,
    ) -> Self {
        crate::init_kernels();
        let dev = Device::Cuda(device);
        assert!(num_stages >= 1);
        let num_layers = config.num_layers();
        assert!(
            num_layers % num_stages == 0,
            "num_layers {num_layers} not divisible by pp {num_stages}"
        );
        let per_stage = num_layers / num_stages;
        let lo = stage * per_stage;
        let hi = lo + per_stage;
        let is_first_stage = stage == 0;
        let is_last_stage = stage == num_stages - 1;

        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) };
        // Pre-transpose like the TP path's `col`/`row` do for world==1 (no shard).
        let wt = |t: Tensor| -> Tensor { t.to_device(dev).transpose(0, 1).contiguous() };
        let placeholder = || Tensor::from_slice(&[bf16::ZERO], &[1, 1]).to_device(dev);

        let embed_tokens = if is_first_stage {
            repl(take(&mut w, "model.embed_tokens.weight"))
        } else {
            placeholder()
        };
        let norm = if is_last_stage {
            repl(take(&mut w, "model.norm.weight"))
        } else {
            placeholder()
        };
        let lm_head_t = if is_last_stage {
            wt(take(&mut w, "lm_head.weight"))
        } else {
            placeholder()
        };

        let rope_cache = RopeCache::new(
            config.max_seq_len(),
            config.head_dim(),
            config.rope_theta.unwrap_or(1_000_000.0) as f32,
        );

        let mut layers = Vec::with_capacity(per_stage);
        eprintln!(
            "[pp] stage {stage}/{num_stages}: layers [{lo}, {hi}) {}{}",
            if is_first_stage { "+embed " } else { "" },
            if is_last_stage { "+norm+lm_head" } else { "" }
        );
        for i in lo..hi {
            let p = format!("model.layers.{i}");
            let q_proj_wt = wt(take(&mut w, &format!("{p}.self_attn.q_proj.weight")));
            let k_proj_wt = wt(take(&mut w, &format!("{p}.self_attn.k_proj.weight")));
            let v_proj_wt = wt(take(&mut w, &format!("{p}.self_attn.v_proj.weight")));
            let q_dim = q_proj_wt.shape()[1];
            let kv_dim = k_proj_wt.shape()[1];
            let qkv_proj_wt = cat_cols(&[&q_proj_wt, &k_proj_wt, &v_proj_wt]);
            drop((q_proj_wt, k_proj_wt, v_proj_wt));
            let gate_proj_wt = wt(take(&mut w, &format!("{p}.mlp.gate_proj.weight")));
            let up_proj_wt = wt(take(&mut w, &format!("{p}.mlp.up_proj.weight")));
            let gate_up_proj_wt = cat_cols(&[&gate_proj_wt, &up_proj_wt]);
            drop((gate_proj_wt, up_proj_wt));
            layers.push(Qwen3Block {
                input_norm: repl(take(&mut w, &format!("{p}.input_layernorm.weight"))),
                qkv_proj_wt,
                q_dim,
                kv_dim,
                o_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
                q_norm: repl(take(&mut w, &format!("{p}.self_attn.q_norm.weight"))),
                k_norm: repl(take(&mut w, &format!("{p}.self_attn.k_norm.weight"))),
                post_norm: repl(take(
                    &mut w,
                    &format!("{p}.post_attention_layernorm.weight"),
                )),
                gate_up_proj_wt,
                down_proj_wt: wt(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
            });
        }

        Self {
            local_num_heads: config.num_heads(),
            local_num_kv_heads: config.num_kv_heads(),
            config,
            embed_tokens,
            layers,
            norm,
            lm_head_t,
            rope_cache,
            tp: None,
            is_first_stage,
            is_last_stage,
        }
    }

    /// Stage-0 token embedding: `[S]` token ids -> `[S, hidden]` hidden state.
    pub fn embed(&self, token_ids: &[u32]) -> Tensor {
        debug_assert!(self.is_first_stage);
        embedding(&self.embed_tokens, token_ids)
    }

    /// Last-stage head: `[*, hidden]` -> logits `[*, vocab]`.
    pub fn head(&self, x: &Tensor) -> Tensor {
        debug_assert!(self.is_last_stage);
        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
        let x = rmsnorm(x, &self.norm, eps);
        matmul_2d(&x, &self.lm_head_t)
    }

    pub fn pp_is_first(&self) -> bool {
        self.is_first_stage
    }
    pub fn pp_is_last(&self) -> bool {
        self.is_last_stage
    }

    /// PP prefill over THIS stage's layers. `x` is `[S, hidden]` (stage 0: from
    /// `embed`; otherwise received from the previous stage). Writes K/V for this
    /// stage's layers into `paged_cache` (indexed by local layer id) and returns
    /// the `[S, hidden]` hidden state to hand to the next stage. Same kernels as
    /// `forward_prefill_paged`, minus embedding and the final norm/lm_head.
    pub fn forward_layers_prefill(
        &self,
        mut x: Tensor,
        slot: usize,
        paged_cache: &mut PagedKVCache,
    ) -> Tensor {
        let new_tokens = x.shape()[0];
        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.config.rms_norm_eps.unwrap_or(1e-6) as f32;

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

        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 = rmsnorm(&x, &layer.input_norm, eps);

            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
            let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
            let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
            let v = qkv
                .narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
                .contiguous();

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

            let q = head_rmsnorm(&q, &layer.q_norm, eps);
            let k = head_rmsnorm(&k, &layer.k_norm, eps);

            let q = xserv_kernels::transpose_for_rope_gpu(&q, new_tokens, num_heads, head_dim);
            let k = xserv_kernels::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 = xserv_kernels::transpose_from_rope_gpu(&q, new_tokens, num_heads, head_dim);
            let k = xserv_kernels::transpose_from_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);

            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);
            let attn_out = flash_attention(&q, &k_full, &v_full, true);

            let attn_merged =
                xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);

            let (normed, x_new) =
                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
            let residual = x_new.clone();

            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
            let ffn_dim = gate_up.shape()[1] / 2;
            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
            x = add_any(&residual, &down);
        }
        x
    }

    /// PP decode over THIS stage's layers. `x` is `[B, hidden]`. Returns
    /// `[B, hidden]`. Positions are read from `paged_cache` (all stages advance
    /// in lockstep, so they agree). Same kernels as `forward_decode_paged`.
    pub fn forward_layers_decode(
        &self,
        mut x: Tensor,
        seq_slots: &[usize],
        paged_cache: &mut PagedKVCache,
    ) -> Tensor {
        let batch = seq_slots.len();
        assert_eq!(x.shape()[0], batch);
        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.config.rms_norm_eps.unwrap_or(1e-6) as f32;

        let positions: Vec<usize> = seq_slots.iter().map(|&s| paged_cache.seq_len(s)).collect();
        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);

        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();

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

            let qkv_all = matmul_2d(&normed, &layer.qkv_proj_wt);
            let q_all = qkv_all.narrow(1, 0, layer.q_dim).contiguous();
            let k_all = qkv_all.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
            let v_all = qkv_all
                .narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
                .contiguous();

            let mut q_rows: Vec<Tensor> = Vec::with_capacity(batch);
            for b in 0..batch {
                let q_row = row_view(&q_all, b);
                let k_row = row_view(&k_all, b);
                let v_row = row_view(&v_all, b);

                let q = xserv_kernels::reshape_heads_gpu(&q_row, 1, num_heads, head_dim);
                let k = xserv_kernels::reshape_heads_gpu(&k_row, 1, num_kv_heads, head_dim);
                let v = xserv_kernels::reshape_heads_gpu(&v_row, 1, num_kv_heads, head_dim);

                let q = head_rmsnorm(&q, &layer.q_norm, eps);
                let k = head_rmsnorm(&k, &layer.k_norm, eps);

                let q = xserv_kernels::transpose_for_rope_gpu(&q, 1, num_heads, head_dim);
                let k = xserv_kernels::transpose_for_rope_gpu(&k, 1, num_kv_heads, head_dim);

                let pos = [positions[b] as u32];
                rope_inplace(&q, &self.rope_cache, &pos);
                rope_inplace(&k, &self.rope_cache, &pos);

                let q = xserv_kernels::transpose_from_rope_gpu(&q, 1, num_heads, head_dim);
                let k = xserv_kernels::transpose_from_rope_gpu(&k, 1, num_kv_heads, head_dim);

                paged_cache.append_tokens(seq_slots[b], layer_idx, &k, &v, 1, positions[b]);

                let q_flat = xserv_kernels::merge_heads_gpu(&q, 1, num_heads, head_dim);
                q_rows.push(q_flat);
            }

            let q_batched_2d = concat_rows(&q_rows);
            let q_4d = q_batched_2d.reshape(&[batch, num_heads, 1, head_dim]);

            let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
            let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;

            let attn_out = xserv_kernels::paged_decode_attention(
                &q_4d,
                k_pool_ptr,
                v_pool_ptr,
                bt_ptr,
                cl_ptr,
                batch,
                num_heads,
                num_kv_heads,
                head_dim,
                max_blocks,
            );

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

            let (normed, x_new) =
                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
            let residual = x_new.clone();

            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
            let ffn_dim = gate_up.shape()[1] / 2;
            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
            x = add_any(&residual, &down);
        }

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

    /// In-place AllReduce(sum) of a partial `[*, hidden]` BF16 activation across
    /// TP ranks (no-op when not tensor-parallel). Used after o_proj and down_proj.
    #[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 std::ffi::c_void;
                tp.all_reduce_sum_bf16_ptr(ptr, t.numel());
            }
        }
    }

    pub fn forward_with_cache(&self, token_ids: &[u32], cache: &mut KVCache) -> Tensor {
        let new_tokens = token_ids.len();
        let pos_offset = cache.seq_len();
        let hidden = self.config.hidden();
        let num_heads = self.config.num_heads();
        let num_kv_heads = self.config.num_kv_heads();
        let head_dim = self.config.head_dim();
        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;

        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 = rmsnorm(&x, &layer.input_norm, eps);

            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
            let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
            let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
            let v = qkv
                .narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
                .contiguous();

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

            let q = head_rmsnorm(&q, &layer.q_norm, eps);
            let k = head_rmsnorm(&k, &layer.k_norm, eps);

            let q = transpose_for_rope(&q, new_tokens, num_heads, head_dim);
            let k = transpose_for_rope(&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(&q, new_tokens, num_heads, head_dim);
            let k = transpose_from_rope(&k, new_tokens, num_kv_heads, head_dim);

            let k_cpu = k.to_device(Device::Cpu);
            let v_cpu = v.to_device(Device::Cpu);
            cache.append_kv_tensor(layer_idx, &k_cpu, &v_cpu, new_tokens);
            let (k_full, v_full) = cache.get_kv_tensors(layer_idx);

            let n_rep = num_heads / num_kv_heads;
            let k_full = repeat_kv(&k_full, n_rep);
            let v_full = repeat_kv(&v_full, n_rep);

            let attn_out = attention(&q, &k_full, &v_full, true);
            let attn_merged = merge_heads_any(&attn_out, new_tokens, hidden);
            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
            x = add_any(&residual, &attn_proj);

            let residual = x.clone();
            let normed = rmsnorm(&x, &layer.post_norm, eps);
            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
            let ffn_dim = gate_up.shape()[1] / 2;
            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
            let gate_activated = silu(&gate);
            let hidden_states = mul_any(&gate_activated, &up);
            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
            x = add_any(&residual, &down);
        }

        let x = rmsnorm(&x, &self.norm, eps);
        matmul_2d(&x, &self.lm_head_t)
    }

    /// Batched decode: process one token per sequence simultaneously.
    /// All compute-heavy ops (projections, FFN) operate on [B, hidden] tensors.
    /// Per-sequence ops (RoPE, KV cache, attention) are handled individually.
    ///
    /// tokens: one token per sequence (len = batch_size)
    /// positions: position offset for each sequence (len = batch_size)
    /// caches: one mutable KV cache per sequence (len = batch_size)
    ///
    /// Returns logits: [batch_size, vocab_size]
    pub fn forward_decode_batch(
        &self,
        tokens: &[u32],
        positions: &[usize],
        caches: &mut [&mut GpuKVCache],
    ) -> Tensor {
        let batch = tokens.len();
        assert_eq!(positions.len(), batch);
        assert_eq!(caches.len(), batch);
        assert!(batch > 0);

        let num_heads = self.config.num_heads();
        let num_kv_heads = self.config.num_kv_heads();
        let head_dim = self.config.head_dim();
        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;

        // Batched embedding: [B, hidden]
        let mut x = embedding(&self.embed_tokens, tokens);

        for (layer_idx, layer) in self.layers.iter().enumerate() {
            let residual = x.clone();
            let normed = rmsnorm(&x, &layer.input_norm, eps); // [B, hidden]

            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
            let q_all = qkv.narrow(1, 0, layer.q_dim).contiguous();
            let k_all = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
            let v_all = qkv
                .narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
                .contiguous();

            // Per-sequence: reshape, qk-norm, RoPE, KV cache, attention, merge
            let mut attn_outputs: Vec<Tensor> = Vec::with_capacity(batch);
            for b in 0..batch {
                // Extract row b: [1, X] — view into contiguous [B, X]
                let q_row = row_view(&q_all, b); // [1, num_heads*head_dim]
                let k_row = row_view(&k_all, b); // [1, num_kv_heads*head_dim]
                let v_row = row_view(&v_all, b); // [1, num_kv_heads*head_dim]

                // GPU reshape: [1, H*D] → [1, H, 1, D]
                let q = xserv_kernels::reshape_heads_gpu(&q_row, 1, num_heads, head_dim);
                let k = xserv_kernels::reshape_heads_gpu(&k_row, 1, num_kv_heads, head_dim);
                let v = xserv_kernels::reshape_heads_gpu(&v_row, 1, num_kv_heads, head_dim);

                // QK norm
                let q = head_rmsnorm(&q, &layer.q_norm, eps);
                let k = head_rmsnorm(&k, &layer.k_norm, eps);

                // GPU transpose for RoPE: [1, H, 1, D] → [1, H, D]
                let q = xserv_kernels::transpose_for_rope_gpu(&q, 1, num_heads, head_dim);
                let k = xserv_kernels::transpose_for_rope_gpu(&k, 1, num_kv_heads, head_dim);

                // RoPE with per-sequence position
                let pos = [positions[b] as u32];
                rope_inplace(&q, &self.rope_cache, &pos);
                rope_inplace(&k, &self.rope_cache, &pos);

                // Transpose back: [1, H, D] → [1, H, 1, D]
                let q = xserv_kernels::transpose_from_rope_gpu(&q, 1, num_heads, head_dim);
                let k = xserv_kernels::transpose_from_rope_gpu(&k, 1, num_kv_heads, head_dim);

                // KV cache: append and get full cache
                let pos_b = positions[b];
                caches[b].append(layer_idx, &k, &v, 1, pos_b);
                let (k_full, v_full) = caches[b].get_kv_len(layer_idx, pos_b + 1);

                // Decode attention (uses native GQA, no repeat_kv needed)
                let attn_out = flash_attention(&q, &k_full, &v_full, true);

                // Merge heads: [1, H, 1, D] → [1, hidden]
                let merged = xserv_kernels::merge_heads_gpu(&attn_out, 1, num_heads, head_dim);
                attn_outputs.push(merged);
            }

            // Concat attention outputs: [B, hidden]
            let attn_merged = concat_rows(&attn_outputs);

            // Batched O projection: [B, hidden] × [hidden, hidden] = [B, hidden]
            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);

            // Fused add + rmsnorm
            let (normed, x_new) =
                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
            let residual = x_new.clone();

            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
            let ffn_dim = gate_up.shape()[1] / 2;
            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
            x = add_any(&residual, &down);
        }

        // Advance KV cache seq_len for each sequence
        for b in 0..batch {
            caches[b].advance_seq_len(1);
        }

        let x = rmsnorm(&x, &self.norm, eps);
        matmul_2d(&x, &self.lm_head_t) // [B, vocab_size]
    }

    /// Paged decode: process one token per sequence using a shared paged KV cache.
    ///
    /// tokens:    [B] one token per sequence
    /// positions: [B] current logical position (BEFORE this step) per sequence
    /// seq_slots: [B] slot ids in `paged_cache`
    ///
    /// Layout note: for S=1 decode the memory of `[B, H, 1, D]`,
    /// `[B, H, D]`, and `[B, H*D]` is the same — only shape/strides differ.
    /// We exploit this to drop every per-sequence kernel: head_rmsnorm and
    /// RoPE both natively accept the `[B*H, D]` / `[B, H, D]` layouts that
    /// fall out of the projection matmuls, and the new-token KV scatter is
    /// one batched `reshape_and_cache` kernel.
    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);

        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 std::ffi::c_void,
            pos_gpu.as_ptr() as *const std::ffi::c_void,
            batch,
            seq_slots,
            paged_cache,
        );
        logits
    }

    /// Host-side per-step cache bookkeeping: block allocation + uploading block
    /// tables / context lens to their (stable-address) GPU buffers. Runs
    /// OUTSIDE any 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);
    }

    /// Pure-GPU decode step: embedding → all 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 makes this region CUDA-graph capturable.
    pub fn decode_core(
        &self,
        ids_gpu: *const std::ffi::c_void,
        pos_gpu: *const std::ffi::c_void,
        batch: usize,
        seq_slots: &[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.config.rms_norm_eps.unwrap_or(1e-6) as f32;

        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 = rmsnorm(&x, &layer.input_norm, eps);

            // Fused QKV projection: one GEMV instead of three.
            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
            let q_dim = num_heads * head_dim;
            let kv_dim = num_kv_heads * head_dim;
            let q_all = qkv.narrow(1, 0, q_dim);
            let k_all = qkv.narrow(1, q_dim, kv_dim);
            let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);

            let q_flat = q_all.contiguous().reshape(&[batch * num_heads, head_dim]);
            let k_flat = k_all
                .contiguous()
                .reshape(&[batch * num_kv_heads, head_dim]);
            let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
            let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);

            let q_3d = q_normed.reshape(&[batch, num_heads, head_dim]);
            let k_3d = k_normed.reshape(&[batch, num_kv_heads, head_dim]);
            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.contiguous().reshape(&[batch, num_kv_heads, head_dim]);

            paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, batch);

            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 std::ffi::c_void;
            let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
            let attn_out = xserv_kernels::paged_decode_attention(
                &q_4d,
                k_pool_ptr,
                v_pool_ptr,
                bt_ptr,
                cl_ptr,
                batch,
                num_heads,
                num_kv_heads,
                head_dim,
                max_blocks,
            );

            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 (normed, x_new) =
                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
            let residual = x_new.clone();

            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
            let ffn_dim = gate_up.shape()[1] / 2;
            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
            self.all_reduce(&down);
            x = add_any(&residual, &down);
        }

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

        let x = rmsnorm(&x, &self.norm, eps);
        matmul_2d(&x, &self.lm_head_t)
    }

    /// Like `decode_core` but also captures hidden states at 3 specified layer
    /// indices (after residual+MLP output). Used by EAGLE3 speculative drafting
    /// to feed the draft head with low/mid/high target representations.
    pub fn decode_core_with_hidden(
        &self,
        ids_gpu: *const std::ffi::c_void,
        pos_gpu: *const std::ffi::c_void,
        batch: usize,
        seq_slots: &[usize],
        paged_cache: &mut PagedKVCache,
        hook_layers: &[usize; 3],
    ) -> (Tensor, [Tensor; 3]) {
        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.config.rms_norm_eps.unwrap_or(1e-6) as f32;

        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);
        let mut hooks: [Option<Tensor>; 3] = [None, None, None];

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

            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
            let q_dim = num_heads * head_dim;
            let kv_dim = num_kv_heads * head_dim;
            let q_all = qkv.narrow(1, 0, q_dim);
            let k_all = qkv.narrow(1, q_dim, kv_dim);
            let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);

            let q_flat = q_all.contiguous().reshape(&[batch * num_heads, head_dim]);
            let k_flat = k_all
                .contiguous()
                .reshape(&[batch * num_kv_heads, head_dim]);
            let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
            let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);

            let q_3d = q_normed.reshape(&[batch, num_heads, head_dim]);
            let k_3d = k_normed.reshape(&[batch, num_kv_heads, head_dim]);
            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.contiguous().reshape(&[batch, num_kv_heads, head_dim]);

            paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, batch);

            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 std::ffi::c_void;
            let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
            let attn_out = xserv_kernels::paged_decode_attention(
                &q_4d,
                k_pool_ptr,
                v_pool_ptr,
                bt_ptr,
                cl_ptr,
                batch,
                num_heads,
                num_kv_heads,
                head_dim,
                max_blocks,
            );

            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 (normed, x_new) =
                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
            let residual = x_new.clone();

            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
            let ffn_dim = gate_up.shape()[1] / 2;
            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
            self.all_reduce(&down);
            x = add_any(&residual, &down);

            for (h_idx, &h_layer) in hook_layers.iter().enumerate() {
                if layer_idx == h_layer {
                    hooks[h_idx] = Some(x.clone());
                }
            }
        }

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

        let x = rmsnorm(&x, &self.norm, eps);
        let logits = matmul_2d(&x, &self.lm_head_t);
        let hidden_arr = [
            hooks[0].take().expect("hook layer 0 not reached"),
            hooks[1].take().expect("hook layer 1 not reached"),
            hooks[2].take().expect("hook layer 2 not reached"),
        ];
        (logits, hidden_arr)
    }

    /// Paged prefill: write a sequence of `new_tokens` K/V into the paged
    /// cache for `slot`, run flash attention via gathered contiguous K/V.
    /// Returns logits [new_tokens, vocab_size].
    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);
        // TP: this rank owns a slice of the heads (local_* == full when world==1).
        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.config.rms_norm_eps.unwrap_or(1e-6) as f32;

        // Pre-allocate enough blocks and bump seq_len up-front so per-layer
        // gather_kv_contiguous returns the freshly written K/V range.
        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 = rmsnorm(&x, &layer.input_norm, eps);

            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
            let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
            let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
            let v = qkv
                .narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
                .contiguous();

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

            let q = head_rmsnorm(&q, &layer.q_norm, eps);
            let k = head_rmsnorm(&k, &layer.k_norm, eps);

            let q = xserv_kernels::transpose_for_rope_gpu(&q, new_tokens, num_heads, head_dim);
            let k = xserv_kernels::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 = xserv_kernels::transpose_from_rope_gpu(&q, new_tokens, num_heads, head_dim);
            let k = xserv_kernels::transpose_from_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);

            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);
            let attn_out = flash_attention(&q, &k_full, &v_full, true);

            let attn_merged =
                xserv_kernels::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 (normed, x_new) =
                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
            let residual = x_new.clone();

            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
            let ffn_dim = gate_up.shape()[1] / 2;
            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
            self.all_reduce(&down);
            x = add_any(&residual, &down);
        }

        let x = rmsnorm(&x, &self.norm, eps);
        matmul_2d(&x, &self.lm_head_t)
    }

    /// Paged multi-token verify path: write `token_ids` into the paged cache,
    /// then verify them with the same paged decode attention kernel used by
    /// single-token decode. This keeps greedy top-1 behavior aligned with
    /// `forward_decode_paged` while still batching the dense projections/MLP
    /// across the draft window.
    pub fn forward_verify_paged_decode_attention(
        &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.config.rms_norm_eps.unwrap_or(1e-6) as f32;

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

        let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens)
            .map(|p| p as u32)
            .collect();
        let kv_lens: Vec<i32> = (0..new_tokens)
            .map(|i| (pos_offset + i + 1) as i32)
            .collect();
        let slots = vec![slot; new_tokens];
        paged_cache.sync_active_batch_with_lens(&slots, &kv_lens);
        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(&self.embed_tokens, token_ids);

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

            let qkv = matmul_batched_gemv(&normed, &layer.qkv_proj_wt);
            let q_dim = num_heads * head_dim;
            let kv_dim = num_kv_heads * head_dim;
            let q_all = qkv.narrow(1, 0, q_dim);
            let k_all = qkv.narrow(1, q_dim, kv_dim);
            let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);

            let q_flat = q_all
                .contiguous()
                .reshape(&[new_tokens * num_heads, head_dim]);
            let k_flat = k_all
                .contiguous()
                .reshape(&[new_tokens * num_kv_heads, head_dim]);
            let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
            let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);

            let q_3d = q_normed.reshape(&[new_tokens, num_heads, head_dim]);
            let k_3d = k_normed.reshape(&[new_tokens, num_kv_heads, head_dim]);
            rope_inplace(&q_3d, &self.rope_cache, &positions);
            rope_inplace(&k_3d, &self.rope_cache, &positions);

            let v_3d = v_all
                .contiguous()
                .reshape(&[new_tokens, num_kv_heads, head_dim]);
            paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, new_tokens);

            let q_decode = q_3d.reshape(&[new_tokens, num_heads, 1, head_dim]);
            let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
            let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
            let attn_out = xserv_kernels::paged_decode_attention(
                &q_decode,
                k_pool_ptr,
                v_pool_ptr,
                bt_ptr,
                cl_ptr,
                new_tokens,
                num_heads,
                num_kv_heads,
                head_dim,
                max_blocks,
            );

            let attn_merged = attn_out.reshape(&[new_tokens, num_heads * head_dim]);
            let attn_proj = matmul_batched_gemv(&attn_merged, &layer.o_proj_wt);
            self.all_reduce(&attn_proj);

            let (normed, x_new) =
                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
            let residual = x_new.clone();

            let gate_up = matmul_batched_gemv(&normed, &layer.gate_up_proj_wt);
            let ffn_dim = gate_up.shape()[1] / 2;
            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
            let down = matmul_batched_gemv(&hidden_states, &layer.down_proj_wt);
            self.all_reduce(&down);
            x = add_any(&residual, &down);
        }

        let x = rmsnorm(&x, &self.norm, eps);
        matmul_batched_gemv(&x, &self.lm_head_t)
    }

    /// Like `forward_verify_paged_decode_attention`, but also captures hidden
    /// states at 3 layer indices (per position). Returns
    /// (logits [new_tokens, vocab], hooks [3][new_tokens, hidden]). Used by
    /// EAGLE3 speculative γ≥2 verify path so we can seed the next round's
    /// EAGLE draft with target's real hidden states at the accepted position.
    pub fn forward_verify_paged_decode_attention_with_hidden(
        &self,
        token_ids: &[u32],
        slot: usize,
        paged_cache: &mut PagedKVCache,
        hook_layers: &[usize; 3],
    ) -> (Tensor, [Tensor; 3]) {
        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.config.rms_norm_eps.unwrap_or(1e-6) as f32;

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

        let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens)
            .map(|p| p as u32)
            .collect();
        let kv_lens: Vec<i32> = (0..new_tokens)
            .map(|i| (pos_offset + i + 1) as i32)
            .collect();
        let slots = vec![slot; new_tokens];
        paged_cache.sync_active_batch_with_lens(&slots, &kv_lens);
        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(&self.embed_tokens, token_ids);
        let mut hooks: [Option<Tensor>; 3] = [None, None, None];

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

            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
            let q_dim = num_heads * head_dim;
            let kv_dim = num_kv_heads * head_dim;
            let q_all = qkv.narrow(1, 0, q_dim);
            let k_all = qkv.narrow(1, q_dim, kv_dim);
            let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);

            let q_flat = q_all
                .contiguous()
                .reshape(&[new_tokens * num_heads, head_dim]);
            let k_flat = k_all
                .contiguous()
                .reshape(&[new_tokens * num_kv_heads, head_dim]);
            let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
            let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);

            let q_3d = q_normed.reshape(&[new_tokens, num_heads, head_dim]);
            let k_3d = k_normed.reshape(&[new_tokens, num_kv_heads, head_dim]);
            rope_inplace(&q_3d, &self.rope_cache, &positions);
            rope_inplace(&k_3d, &self.rope_cache, &positions);

            let v_3d = v_all
                .contiguous()
                .reshape(&[new_tokens, num_kv_heads, head_dim]);
            paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, new_tokens);

            let q_decode = q_3d.reshape(&[new_tokens, num_heads, 1, head_dim]);
            let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
            let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
            let attn_out = xserv_kernels::paged_decode_attention(
                &q_decode,
                k_pool_ptr,
                v_pool_ptr,
                bt_ptr,
                cl_ptr,
                new_tokens,
                num_heads,
                num_kv_heads,
                head_dim,
                max_blocks,
            );

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

            let (normed, x_new) =
                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
            let residual = x_new.clone();

            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
            let ffn_dim = gate_up.shape()[1] / 2;
            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
            self.all_reduce(&down);
            x = add_any(&residual, &down);

            for (h_idx, &h_layer) in hook_layers.iter().enumerate() {
                if layer_idx == h_layer {
                    hooks[h_idx] = Some(x.clone());
                }
            }
        }

        let x = rmsnorm(&x, &self.norm, eps);
        let logits = matmul_2d(&x, &self.lm_head_t);
        let hidden_arr = [
            hooks[0].take().expect("hook layer 0 not reached"),
            hooks[1].take().expect("hook layer 1 not reached"),
            hooks[2].take().expect("hook layer 2 not reached"),
        ];
        (logits, hidden_arr)
    }

    /// Tree-aware verify: like `_with_hidden` but supports sibling candidates
    /// sharing the same target position. Caller supplies per-token positions
    /// (for RoPE), kv_lens (attention context length), and a flattened
    /// `tree_mask` (`[new_tokens, new_tokens]` i32; `mask[i, j]!=0` iff query i
    /// attends to newly-written K/V at slot j). Positions in the paged cache
    /// before pos_offset are always attended (regular history).
    #[allow(clippy::too_many_arguments)]
    pub fn forward_verify_paged_decode_attention_tree_with_hidden(
        &self,
        token_ids: &[u32],
        positions: &[u32],
        kv_lens: &[i32],
        tree_mask: &[i32],
        slot: usize,
        paged_cache: &mut PagedKVCache,
        hook_layers: &[usize; 3],
    ) -> (Tensor, [Tensor; 3]) {
        let new_tokens = token_ids.len();
        assert_eq!(positions.len(), new_tokens);
        assert_eq!(kv_lens.len(), new_tokens);
        assert_eq!(tree_mask.len(), new_tokens * new_tokens);

        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.config.rms_norm_eps.unwrap_or(1e-6) as f32;

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

        let slots = vec![slot; new_tokens];
        paged_cache.sync_active_batch_with_lens(&slots, kv_lens);
        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();

        // Upload tree_mask [new_tokens, new_tokens] i32 to GPU.
        let mask_bytes: &[u8] = unsafe {
            std::slice::from_raw_parts(tree_mask.as_ptr() as *const u8, tree_mask.len() * 4)
        };
        let mut mask_buf =
            xserv_cuda::allocator::cached_alloc(mask_bytes.len()).expect("alloc tree_mask");
        mask_buf.copy_from_host(mask_bytes).unwrap();
        let mask_ptr = mask_buf.as_ptr() as *const i32;

        let mut x = embedding(&self.embed_tokens, token_ids);
        let mut hooks: [Option<Tensor>; 3] = [None, None, None];

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

            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
            let q_dim = num_heads * head_dim;
            let kv_dim = num_kv_heads * head_dim;
            let q_all = qkv.narrow(1, 0, q_dim);
            let k_all = qkv.narrow(1, q_dim, kv_dim);
            let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);

            let q_flat = q_all
                .contiguous()
                .reshape(&[new_tokens * num_heads, head_dim]);
            let k_flat = k_all
                .contiguous()
                .reshape(&[new_tokens * num_kv_heads, head_dim]);
            let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
            let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);

            let q_3d = q_normed.reshape(&[new_tokens, num_heads, head_dim]);
            let k_3d = k_normed.reshape(&[new_tokens, num_kv_heads, head_dim]);
            rope_inplace(&q_3d, &self.rope_cache, positions);
            rope_inplace(&k_3d, &self.rope_cache, positions);

            let v_3d = v_all
                .contiguous()
                .reshape(&[new_tokens, num_kv_heads, head_dim]);
            paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, new_tokens);

            let q_decode = q_3d.reshape(&[new_tokens, num_heads, 1, head_dim]);
            let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
            let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
            let attn_out = xserv_kernels::paged_decode_attention_tree(
                &q_decode,
                k_pool_ptr,
                v_pool_ptr,
                bt_ptr,
                cl_ptr,
                mask_ptr,
                new_tokens,
                num_heads,
                num_kv_heads,
                head_dim,
                max_blocks,
                pos_offset,
                new_tokens,
            );

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

            let (normed, x_new) =
                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
            let residual = x_new.clone();

            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
            let ffn_dim = gate_up.shape()[1] / 2;
            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
            self.all_reduce(&down);
            x = add_any(&residual, &down);

            for (h_idx, &h_layer) in hook_layers.iter().enumerate() {
                if layer_idx == h_layer {
                    hooks[h_idx] = Some(x.clone());
                }
            }
        }

        let x = rmsnorm(&x, &self.norm, eps);
        let logits = matmul_2d(&x, &self.lm_head_t);
        let hidden_arr = [
            hooks[0].take().expect("hook layer 0 not reached"),
            hooks[1].take().expect("hook layer 1 not reached"),
            hooks[2].take().expect("hook layer 2 not reached"),
        ];
        (logits, hidden_arr)
    }

    /// Forward with GPU-resident KV cache and GPU transpose/reshape kernels.
    pub fn forward_gpu_cache(&self, token_ids: &[u32], cache: &mut GpuKVCache) -> Tensor {
        let new_tokens = token_ids.len();
        let pos_offset = cache.seq_len();

        let num_heads = self.config.num_heads();
        let num_kv_heads = self.config.num_kv_heads();
        let head_dim = self.config.head_dim();
        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;

        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 = rmsnorm(&x, &layer.input_norm, eps);

            let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
            let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
            let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
            let v = qkv
                .narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
                .contiguous();

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

            let q = head_rmsnorm(&q, &layer.q_norm, eps);
            let k = head_rmsnorm(&k, &layer.k_norm, eps);

            let q = xserv_kernels::transpose_for_rope_gpu(&q, new_tokens, num_heads, head_dim);
            let k = xserv_kernels::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 = xserv_kernels::transpose_from_rope_gpu(&q, new_tokens, num_heads, head_dim);
            let k = xserv_kernels::transpose_from_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);

            cache.append(layer_idx, &k, &v, new_tokens, pos_offset);
            let (k_full, v_full) = cache.get_kv_len(layer_idx, pos_offset + new_tokens);

            let attn_out = flash_attention(&q, &k_full, &v_full, true);
            let attn_merged =
                xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);

            let (normed, x_new) =
                xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
            let residual = x_new.clone();

            let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
            let ffn_dim = gate_up.shape()[1] / 2;
            let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
            let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
            x = add_any(&residual, &down);
        }

        cache.advance_seq_len(new_tokens);
        let x = rmsnorm(&x, &self.norm, eps);
        matmul_2d(&x, &self.lm_head_t)
    }

    /// Reference to the target's token embedding table. Shared (not copied)
    /// with speculative draft heads like EAGLE3.
    pub fn embed_tokens_tensor(&self) -> &Tensor {
        &self.embed_tokens
    }

    /// Extract weight pointers for CUDA Graph capture.
    pub fn layer_weight_ptrs(&self) -> Vec<crate::decode_graph::LayerWeightPtrs> {
        self.layers
            .iter()
            .map(|l| crate::decode_graph::LayerWeightPtrs {
                input_norm: l.input_norm.data_ptr() as *const std::ffi::c_void,
                q_proj_wt: l.q_proj_wt().data_ptr() as *const std::ffi::c_void,
                k_proj_wt: l.k_proj_wt().data_ptr() as *const std::ffi::c_void,
                v_proj_wt: l.v_proj_wt().data_ptr() as *const std::ffi::c_void,
                o_proj_wt: l.o_proj_wt.data_ptr() as *const std::ffi::c_void,
                q_norm: l.q_norm.data_ptr() as *const std::ffi::c_void,
                k_norm: l.k_norm.data_ptr() as *const std::ffi::c_void,
                post_norm: l.post_norm.data_ptr() as *const std::ffi::c_void,
                gate_proj_wt: l.gate_proj_wt().data_ptr() as *const std::ffi::c_void,
                up_proj_wt: l.up_proj_wt().data_ptr() as *const std::ffi::c_void,
                down_proj_wt: l.down_proj_wt.data_ptr() as *const std::ffi::c_void,
            })
            .collect()
    }

    /// Get pointers needed for CUDA Graph capture.
    pub fn graph_capture_ptrs(
        &self,
    ) -> (
        *const std::ffi::c_void, // norm weight
        *const std::ffi::c_void, // lm_head_t
        *const std::ffi::c_void, // embed_tokens
        *const std::ffi::c_void, // rope cos
        *const std::ffi::c_void, // rope sin
    ) {
        (
            self.norm.data_ptr() as *const std::ffi::c_void,
            self.lm_head_t.data_ptr() as *const std::ffi::c_void,
            self.embed_tokens.data_ptr() as *const std::ffi::c_void,
            self.rope_cache.cos.as_ptr() as *const std::ffi::c_void,
            self.rope_cache.sin.as_ptr() as *const std::ffi::c_void,
        )
    }
}

// --- Helpers ---

/// Keep this rank's contiguous row-block of a 2D `[rows, cols]` BF16 tensor
/// (column-parallel split: split the OUTPUT dim). `world==1` returns the whole.
/// Input must be a contiguous CPU (or device) BF16 tensor.
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, "shard_rows expects 2D weight");
    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])
}

/// Keep this rank's column-block of a 2D `[rows, cols]` BF16 tensor (row-parallel
/// split: split the INPUT dim). Strided copy. `world==1` returns the whole.
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, "shard_cols expects 2D weight");
    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 matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
    assert_eq!(a.ndim(), 2);
    assert_eq!(b.ndim(), 2);
    matmul(a, b, GemmBackend::CuBlas)
}

fn reshape_heads(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
    let x_cpu = x.to_device(Device::Cpu);
    let hidden = num_heads * head_dim;
    let src = x_cpu.as_slice::<bf16>();
    let mut out = vec![bf16::ZERO; num_heads * seq_len * head_dim];
    for s in 0..seq_len {
        for h in 0..num_heads {
            let si = s * hidden + h * head_dim;
            let di = (h * seq_len + s) * head_dim;
            out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
        }
    }
    Tensor::from_slice(&out, &[1, num_heads, seq_len, head_dim]).to_device(x.device())
}

fn merge_heads_any(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
    let num_heads = x.shape()[1];
    let head_dim = x.shape()[3];
    let x_cpu = x.to_device(Device::Cpu);
    let src = x_cpu.as_slice::<bf16>();
    let mut out = vec![bf16::ZERO; seq_len * hidden];
    for s in 0..seq_len {
        for h in 0..num_heads {
            let si = (h * seq_len + s) * head_dim;
            let di = s * hidden + h * head_dim;
            out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
        }
    }
    Tensor::from_slice(&out, &[seq_len, hidden]).to_device(x.device())
}

/// Per-head RMSNorm: apply RMSNorm to each [head_dim] slice independently.
/// x: [1, H, S, D], norm_weight: [D]
fn head_rmsnorm(x: &Tensor, norm_weight: &Tensor, eps: f32) -> Tensor {
    let num_heads = x.shape()[1];
    let seq_len = x.shape()[2];
    let head_dim = x.shape()[3];
    // Reshape to [H*S, D], apply rmsnorm, reshape back
    let total_rows = num_heads * seq_len;
    let flat = x.reshape(&[total_rows, head_dim]);
    let normed = rmsnorm(&flat, norm_weight, eps);
    normed.reshape(&[1, num_heads, seq_len, head_dim])
}

/// [1, H, S, D] → [S, H, D] for RoPE kernel
fn transpose_for_rope(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
    let x_cpu = x.to_device(Device::Cpu);
    let src = x_cpu.as_slice::<bf16>();
    let mut out = vec![bf16::ZERO; seq_len * num_heads * head_dim];
    for h in 0..num_heads {
        for s in 0..seq_len {
            let si = (h * seq_len + s) * head_dim;
            let di = (s * num_heads + h) * head_dim;
            out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
        }
    }
    Tensor::from_slice(&out, &[seq_len, num_heads, head_dim]).to_device(x.device())
}

/// [S, H, D] → [1, H, S, D] after RoPE
fn transpose_from_rope(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
    let x_cpu = x.to_device(Device::Cpu);
    let src = x_cpu.as_slice::<bf16>();
    let mut out = vec![bf16::ZERO; num_heads * seq_len * head_dim];
    for s in 0..seq_len {
        for h in 0..num_heads {
            let si = (s * num_heads + h) * head_dim;
            let di = (h * seq_len + s) * head_dim;
            out[di..di + head_dim].copy_from_slice(&src[si..si + head_dim]);
        }
    }
    Tensor::from_slice(&out, &[1, num_heads, seq_len, head_dim]).to_device(x.device())
}

fn repeat_kv(x: &Tensor, n_rep: usize) -> Tensor {
    if n_rep == 1 {
        return x.clone();
    }
    let kv_heads = x.shape()[1];
    let seq_len = x.shape()[2];
    let head_dim = x.shape()[3];
    let x_cpu = x.to_device(Device::Cpu);
    let src = x_cpu.as_slice::<bf16>();
    let new_heads = kv_heads * n_rep;
    let mut out = vec![bf16::ZERO; new_heads * seq_len * head_dim];
    let chunk = seq_len * head_dim;
    for kv_h in 0..kv_heads {
        for r in 0..n_rep {
            let dst_h = kv_h * n_rep + r;
            out[dst_h * chunk..(dst_h + 1) * chunk]
                .copy_from_slice(&src[kv_h * chunk..(kv_h + 1) * chunk]);
        }
    }
    Tensor::from_slice(&out, &[1, new_heads, seq_len, head_dim]).to_device(x.device())
}

/// Extract row `b` from a contiguous 2D tensor [B, cols] as a [1, cols] view.
/// Zero-copy: shares storage with the original tensor.
fn row_view(t: &Tensor, row: usize) -> Tensor {
    assert_eq!(t.ndim(), 2);
    assert!(t.is_contiguous());
    let cols = t.shape()[1];
    assert!(row < t.shape()[0]);
    let new_offset = t.offset() + row * cols;
    Tensor::from_storage(
        t.storage().clone(),
        smallvec::SmallVec::from_slice(&[1, cols]),
        xserv_tensor::shape::contiguous_strides(&[1, cols]),
        new_offset,
        t.dtype(),
    )
}

/// 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
}

/// Concatenate row tensors [1, cols] into a single [B, cols] tensor via D2D memcpy.
fn concat_rows(rows: &[Tensor]) -> Tensor {
    assert!(!rows.is_empty());
    let batch = rows.len();
    let cols = rows[0].shape()[1];
    let dtype = rows[0].dtype();
    let device = rows[0].device();
    let elem_size = dtype.size_bytes();
    let row_bytes = cols * elem_size;

    // Allocate output [B, cols] and copy each row into it
    let total_bytes = batch * row_bytes;
    let mut out_buf = xserv_cuda::allocator::cached_alloc(total_bytes).expect("alloc concat_rows");

    for (b, row) in rows.iter().enumerate() {
        assert_eq!(row.shape(), &[1, cols]);
        assert!(row.is_contiguous());
        let src_buf = row.storage().gpu_buffer();
        let src_offset = row.offset() * elem_size;
        let dst_offset = b * row_bytes;
        out_buf
            .copy_from_device_at(src_buf, src_offset, dst_offset, row_bytes)
            .unwrap();
    }

    // Wrap in a Tensor
    let device_id = match device {
        Device::Cuda(id) => id,
        _ => panic!("expected CUDA device"),
    };
    unsafe {
        crate::kv_cache::tensor_from_gpu_buffer_pub(out_buf, &[batch, cols], dtype, device_id)
    }
}

/// Concatenate 2D GPU tensors along dim=1 (columns). All must share dim 0.
fn cat_cols(tensors: &[&Tensor]) -> Tensor {
    assert!(!tensors.is_empty());
    let rows = tensors[0].shape()[0];
    let dtype = tensors[0].dtype();
    let device = tensors[0].device();
    let elem = dtype.size_bytes();
    let total_cols: usize = tensors
        .iter()
        .map(|t| {
            assert_eq!(t.ndim(), 2);
            assert_eq!(t.shape()[0], rows);
            assert!(t.is_contiguous());
            t.shape()[1]
        })
        .sum();
    let out = Tensor::empty(&[rows, total_cols], dtype, device);
    let dst_base = out.data_ptr() as *mut u8;
    for r in 0..rows {
        let mut col_off = 0usize;
        for t in tensors {
            let cols = t.shape()[1];
            let src = unsafe { t.data_ptr().add(r * cols * elem) };
            let dst = unsafe { dst_base.add((r * total_cols + col_off) * elem) };
            let count = cols * elem;
            unsafe {
                xserv_cuda::ffi::cudaMemcpy(
                    dst as *mut u8,
                    src as *const u8,
                    count,
                    2, // cudaMemcpyDeviceToDevice
                );
            }
            col_off += cols;
        }
    }
    out
}

fn add_any(a: &Tensor, b: &Tensor) -> Tensor {
    xserv_kernels::add(a, b)
}

fn mul_any(a: &Tensor, b: &Tensor) -> Tensor {
    xserv_kernels::mul(a, b)
}

pub fn sample_greedy(logits: &Tensor) -> u32 {
    assert_eq!(logits.ndim(), 2);
    let logits_cpu = logits.to_device(Device::Cpu);
    let vocab_size = logits.shape()[1];
    let seq_len = logits.shape()[0];
    let data = logits_cpu.as_slice::<bf16>();
    let last = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
    last.iter()
        .enumerate()
        .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
        .map(|(i, _)| i as u32)
        .unwrap()
}