from __future__ import annotations

import math
import os
from functools import lru_cache

import torch
import torch.nn as nn


HIDDEN = 1024
INTERMEDIATE = 3072
NUM_Q = 16
NUM_KV = 8
HEAD_DIM = 128
NUM_LAYERS = 4
EPS = 1.0e-6


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.input_ln = nn.Parameter(torch.ones(HIDDEN, dtype=torch.bfloat16))
        self.q_proj = nn.Parameter(torch.empty(NUM_Q * HEAD_DIM, HIDDEN, dtype=torch.bfloat16))
        self.k_proj = nn.Parameter(torch.empty(NUM_KV * HEAD_DIM, HIDDEN, dtype=torch.bfloat16))
        self.v_proj = nn.Parameter(torch.empty(NUM_KV * HEAD_DIM, HIDDEN, dtype=torch.bfloat16))
        self.q_norm = nn.Parameter(torch.ones(HEAD_DIM, dtype=torch.bfloat16))
        self.k_norm = nn.Parameter(torch.ones(HEAD_DIM, dtype=torch.bfloat16))
        self.o_proj = nn.Parameter(torch.empty(HIDDEN, NUM_Q * HEAD_DIM, dtype=torch.bfloat16))
        self.post_ln = nn.Parameter(torch.ones(HIDDEN, dtype=torch.bfloat16))
        self.gate_proj = nn.Parameter(torch.empty(INTERMEDIATE, HIDDEN, dtype=torch.bfloat16))
        self.up_proj = nn.Parameter(torch.empty(INTERMEDIATE, HIDDEN, dtype=torch.bfloat16))
        self.down_proj = nn.Parameter(torch.empty(HIDDEN, INTERMEDIATE, dtype=torch.bfloat16))
        for name, p in self.named_parameters():
            if name not in ("input_ln", "q_norm", "k_norm", "post_ln"):
                nn.init.normal_(p, std=0.02)


class Model(nn.Module):
    def __init__(self, num_layers: int = NUM_LAYERS, max_seq: int = 131072):
        super().__init__()
        self.num_layers = num_layers
        self.max_seq = max_seq
        self.blocks = nn.ModuleList([Block() for _ in range(num_layers)])


_CPP_SOURCE = r"""
#include <torch/extension.h>
#include <vector>

void megaqwen_decode_cuda(
    at::Tensor input, at::Tensor noise, at::Tensor h0, at::Tensor h1,
    std::vector<at::Tensor> weights,
    std::vector<at::Tensor> k_caches, std::vector<at::Tensor> v_caches,
    at::Tensor rope_cos, at::Tensor rope_sin,
    at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor attn,
    at::Tensor activations, at::Tensor residual, at::Tensor mlp,
    at::Tensor partial_max, at::Tensor partial_sum, at::Tensor partial_out,
    int64_t num_layers, int64_t pos, int64_t max_seq, int64_t max_chunks);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("decode", &megaqwen_decode_cuda, "Qwen3 decode cooperative CUDA kernel");
}
"""


_CUDA_SOURCE = r"""
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
#include <algorithm>
#include <cfloat>
#include <vector>

namespace cg = cooperative_groups;
constexpr int H = 1024;
constexpr int I = 3072;
constexpr int NQ = 16;
constexpr int NK = 8;
constexpr int D = 128;
constexpr int T = 256;
constexpr int W = T / 32;
constexpr int CHUNK = 256;
constexpr int MAX_LAYERS = 8;

struct LayerWeights {
    const __nv_bfloat16* input_ln;
    const __nv_bfloat16* q_proj;
    const __nv_bfloat16* k_proj;
    const __nv_bfloat16* v_proj;
    const __nv_bfloat16* q_norm;
    const __nv_bfloat16* k_norm;
    const __nv_bfloat16* o_proj;
    const __nv_bfloat16* post_ln;
    const __nv_bfloat16* gate_proj;
    const __nv_bfloat16* up_proj;
    const __nv_bfloat16* down_proj;
};

struct KernelArgs {
    const __nv_bfloat16* input;
    const __nv_bfloat16* noise;
    __nv_bfloat16* h0;
    __nv_bfloat16* h1;
    LayerWeights weights[MAX_LAYERS];
    __nv_bfloat16* k_cache[MAX_LAYERS];
    __nv_bfloat16* v_cache[MAX_LAYERS];
    const float* rope_cos;
    const float* rope_sin;
    float* q;
    float* k;
    float* v;
    float* attn;
    float* activations;
    float* residual;
    float* mlp;
    float* partial_max;
    float* partial_sum;
    float* partial_out;
    int num_layers;
    int pos;
    int max_seq;
    int max_chunks;
};

__device__ __forceinline__ float bf(__nv_bfloat16 x) {
    return __bfloat162float(x);
}

__device__ __forceinline__ float warp_sum(float x) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        x += __shfl_down_sync(0xffffffff, x, off);
    return x;
}

__device__ __forceinline__ float block_sum(float x, float* tmp) {
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    x = warp_sum(x);
    if (lane == 0) tmp[warp] = x;
    __syncthreads();
    if (warp == 0) {
        float v = lane < W ? tmp[lane] : 0.0f;
        v = warp_sum(v);
        if (lane == 0) tmp[0] = v;
    }
    __syncthreads();
    return tmp[0];
}

__device__ __forceinline__ float warp_max(float x) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        x = fmaxf(x, __shfl_down_sync(0xffffffff, x, off));
    return x;
}

__device__ __forceinline__ float block_max(float x, float* tmp) {
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    x = warp_max(x);
    if (lane == 0) tmp[warp] = x;
    __syncthreads();
    if (warp == 0) {
        float v = lane < W ? tmp[lane] : -FLT_MAX;
        v = warp_max(v);
        if (lane == 0) tmp[0] = v;
    }
    __syncthreads();
    return tmp[0];
}

__device__ __forceinline__ float row_dot(
    const __nv_bfloat16* __restrict__ w, const float* __restrict__ x, int n) {
    int lane = threadIdx.x & 31;
    float sum = 0.0f;
    #pragma unroll 4
    for (int k = lane * 8; k < n; k += 256) {
        uint4 bits = __ldg(reinterpret_cast<const uint4*>(w + k));
        const __nv_bfloat16* p = reinterpret_cast<const __nv_bfloat16*>(&bits);
        sum += bf(p[0]) * x[k] + bf(p[1]) * x[k + 1]
             + bf(p[2]) * x[k + 2] + bf(p[3]) * x[k + 3];
        sum += bf(p[4]) * x[k + 4] + bf(p[5]) * x[k + 5]
             + bf(p[6]) * x[k + 6] + bf(p[7]) * x[k + 7];
    }
    return warp_sum(sum);
}

__device__ void qkv_project(
    cg::grid_group& grid, const KernelArgs& a, const LayerWeights& w,
    const __nv_bfloat16* src, const __nv_bfloat16* noise,
    float* q, float* k, float* v, float* residual, bool mix_input, bool wait_input) {
    if (wait_input) grid.sync();
    __shared__ float xnorm[H];
    __shared__ float red[W];
    __shared__ float inv_rms;
    float local = 0.0f;
    for (int i = threadIdx.x; i < H; i += T) {
        float x = bf(src[i]);
        if (mix_input) x = 0.5f * x + 0.5f * bf(noise[i]);
        xnorm[i] = x;
        if (blockIdx.x == 0) residual[i] = x;
        local += x * x;
    }
    float s = warp_sum(local);
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    if (lane == 0) red[warp] = s;
    __syncthreads();
    if (warp == 0) {
        float z = lane < W ? red[lane] : 0.0f;
        z = warp_sum(z);
        if (lane == 0) inv_rms = rsqrtf(z / float(H) + 1.0e-6f);
    }
    __syncthreads();
    for (int i = threadIdx.x; i < H; i += T)
        xnorm[i] = xnorm[i] * inv_rms * bf(w.input_ln[i]);
    __syncthreads();

    constexpr int TOTAL = NQ * D + 2 * NK * D;
    int rows_per_block = (TOTAL + gridDim.x - 1) / gridDim.x;
    int begin = int(blockIdx.x) * rows_per_block;
    int end = min(begin + rows_per_block, TOTAL);
    for (int row = begin + warp; row < end; row += W) {
        const __nv_bfloat16* wr;
        float* out;
        if (row < NQ * D) {
            wr = w.q_proj + row * H; out = q + row;
        } else if (row < NQ * D + NK * D) {
            int r = row - NQ * D; wr = w.k_proj + r * H; out = k + r;
        } else {
            int r = row - NQ * D - NK * D; wr = w.v_proj + r * H; out = v + r;
        }
        float value = row_dot(wr, xnorm, H);
        if (lane == 0) *out = value;
    }
}

__device__ __forceinline__ uint32_t pack_bf16_pair(float a, float b) {
    union Bits { __nv_bfloat16 bf; uint16_t u; };
    Bits lo, hi;
    lo.bf = __float2bfloat16_rn(a);
    hi.bf = __float2bfloat16_rn(b);
    return uint32_t(lo.u) | (uint32_t(hi.u) << 16);
}

__device__ void qk_norm_rope(
    cg::grid_group& grid, const KernelArgs& a, const LayerWeights& w,
    float* q, float* k, float* v, __nv_bfloat16* kc, __nv_bfloat16* vc) {
    grid.sync();
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int qh_per_block = (NQ + gridDim.x - 1) / gridDim.x;
    int q0 = int(blockIdx.x) * qh_per_block;
    int q1 = min(q0 + qh_per_block, NQ);
    for (int h = q0 + warp; h < q1; h += W) {
        int base = lane * 4;
        float* qbase = q + h * D + base;
        float4 x = *reinterpret_cast<const float4*>(qbase);
        uint2 wbits = __ldg(reinterpret_cast<const uint2*>(w.q_norm + base));
        const __nv_bfloat16* nw = reinterpret_cast<const __nv_bfloat16*>(&wbits);
        float ss = 0.0f;
        ss = x.x * x.x + x.y * x.y + x.z * x.z + x.w * x.w;
        ss = warp_sum(ss);
        float inv = rsqrtf(__shfl_sync(0xffffffff, ss / float(D) + 1.0e-6f, 0));
        float l0 = x.x * inv * bf(nw[0]);
        float l1 = x.y * inv * bf(nw[1]);
        float l2 = x.z * inv * bf(nw[2]);
        float l3 = x.w * inv * bf(nw[3]);
        int peer = lane < 16 ? lane + 16 : lane - 16;
        float p0 = __shfl_sync(0xffffffff, l0, peer);
        float p1 = __shfl_sync(0xffffffff, l1, peer);
        float p2 = __shfl_sync(0xffffffff, l2, peer);
        float p3 = __shfl_sync(0xffffffff, l3, peer);
        int fbase = lane < 16 ? base : base - D / 2;
        const float* cbase = a.rope_cos + a.pos * (D / 2) + fbase;
        const float* sbase = a.rope_sin + a.pos * (D / 2) + fbase;
        float c0 = cbase[0], c1 = cbase[1], c2 = cbase[2], c3 = cbase[3];
        float s0 = sbase[0], s1 = sbase[1], s2 = sbase[2], s3 = sbase[3];
        float4 out;
        if (lane < 16) {
            out = make_float4(l0 * c0 - p0 * s0, l1 * c1 - p1 * s1,
                              l2 * c2 - p2 * s2, l3 * c3 - p3 * s3);
        } else {
            out = make_float4(p0 * s0 + l0 * c0, p1 * s1 + l1 * c1,
                              p2 * s2 + l2 * c2, p3 * s3 + l3 * c3);
        }
        *reinterpret_cast<float4*>(qbase) = out;
    }
    int kh_per_block = (NK + gridDim.x - 1) / gridDim.x;
    int k0 = int(blockIdx.x) * kh_per_block;
    int k1 = min(k0 + kh_per_block, NK);
    for (int h = k0 + warp; h < k1; h += W) {
        int base = lane * 4;
        float* kbase = k + h * D + base;
        float4 x = *reinterpret_cast<const float4*>(kbase);
        uint2 wbits = __ldg(reinterpret_cast<const uint2*>(w.k_norm + base));
        const __nv_bfloat16* nw = reinterpret_cast<const __nv_bfloat16*>(&wbits);
        float ss = 0.0f;
        ss = x.x * x.x + x.y * x.y + x.z * x.z + x.w * x.w;
        ss = warp_sum(ss);
        float inv = rsqrtf(__shfl_sync(0xffffffff, ss / float(D) + 1.0e-6f, 0));
        float l0 = x.x * inv * bf(nw[0]);
        float l1 = x.y * inv * bf(nw[1]);
        float l2 = x.z * inv * bf(nw[2]);
        float l3 = x.w * inv * bf(nw[3]);
        int peer = lane < 16 ? lane + 16 : lane - 16;
        float p0 = __shfl_sync(0xffffffff, l0, peer);
        float p1 = __shfl_sync(0xffffffff, l1, peer);
        float p2 = __shfl_sync(0xffffffff, l2, peer);
        float p3 = __shfl_sync(0xffffffff, l3, peer);
        int fbase = lane < 16 ? base : base - D / 2;
        const float* cbase = a.rope_cos + a.pos * (D / 2) + fbase;
        const float* sbase = a.rope_sin + a.pos * (D / 2) + fbase;
        float c0 = cbase[0], c1 = cbase[1], c2 = cbase[2], c3 = cbase[3];
        float s0 = sbase[0], s1 = sbase[1], s2 = sbase[2], s3 = sbase[3];
        float4 out;
        if (lane < 16) {
            out = make_float4(l0 * c0 - p0 * s0, l1 * c1 - p1 * s1,
                              l2 * c2 - p2 * s2, l3 * c3 - p3 * s3);
        } else {
            out = make_float4(p0 * s0 + l0 * c0, p1 * s1 + l1 * c1,
                              p2 * s2 + l2 * c2, p3 * s3 + l3 * c3);
        }
        *reinterpret_cast<float4*>(kbase) = out;
        __nv_bfloat16* kpos = kc + h * a.max_seq * D + a.pos * D + base;
        __nv_bfloat16* vpos = vc + h * a.max_seq * D + a.pos * D + base;
        *reinterpret_cast<uint2*>(kpos) =
            make_uint2(pack_bf16_pair(out.x, out.y), pack_bf16_pair(out.z, out.w));
        const float* vbase = v + h * D + base;
        float4 vv = *reinterpret_cast<const float4*>(vbase);
        *reinterpret_cast<uint2*>(vpos) =
            make_uint2(pack_bf16_pair(vv.x, vv.y), pack_bf16_pair(vv.z, vv.w));
    }
}

__device__ void attention_partials(
    cg::grid_group& grid, const KernelArgs& a, const float* q,
    const __nv_bfloat16* kc, const __nv_bfloat16* vc, float* partial_max,
    float* partial_sum, float* partial_out) {
    grid.sync();
    __shared__ float scores[2][CHUNK];
    __shared__ float reduce_tmp[W];
    __shared__ float maxima[2];
    __shared__ float sums[2];
    __shared__ float acc[2][W][4][32];
    int chunks = (a.pos + 1 + CHUNK - 1) / CHUNK;
    int tasks = NK * chunks;
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int tid = threadIdx.x;
    constexpr float scale = 0.08838834764831845f;
    for (int task = int(blockIdx.x); task < tasks; task += int(gridDim.x)) {
        int kvh = task / chunks;
        int chunk = task - kvh * chunks;
        int begin = chunk * CHUNK;
        int end = min(begin + CHUNK, a.pos + 1);
        int len = end - begin;
        int qh0 = kvh * 2;
        int qh1 = qh0 + 1;
        float q0[4], q1[4];
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            int d = lane * 4 + j;
            q0[j] = q[qh0 * D + d];
            q1[j] = q[qh1 * D + d];
        }
        for (int ix = warp; ix < len; ix += W) {
            const __nv_bfloat16* kp = kc + kvh * a.max_seq * D + (begin + ix) * D;
            uint2 bits = __ldg(reinterpret_cast<const uint2*>(kp + lane * 4));
            const __nv_bfloat16* x = reinterpret_cast<const __nv_bfloat16*>(&bits);
            float z0 = q0[0] * bf(x[0]) + q0[1] * bf(x[1])
                     + q0[2] * bf(x[2]) + q0[3] * bf(x[3]);
            float z1 = q1[0] * bf(x[0]) + q1[1] * bf(x[1])
                     + q1[2] * bf(x[2]) + q1[3] * bf(x[3]);
            z0 = warp_sum(z0) * scale;
            z1 = warp_sum(z1) * scale;
            if (lane == 0) {
                scores[0][ix] = z0;
                scores[1][ix] = z1;
            }
        }
        __syncthreads();
        if (tid == 0) {
            float m0 = -FLT_MAX, m1 = -FLT_MAX;
            for (int i = 0; i < len; ++i) {
                m0 = fmaxf(m0, scores[0][i]);
                m1 = fmaxf(m1, scores[1][i]);
            }
            maxima[0] = m0;
            maxima[1] = m1;
        }
        __syncthreads();
        float s0 = 0.0f, s1 = 0.0f;
        for (int i = tid; i < len; i += T) {
            s0 += expf(scores[0][i] - maxima[0]);
            s1 += expf(scores[1][i] - maxima[1]);
        }
        float sum0 = block_sum(s0, reduce_tmp);
        float sum1 = block_sum(s1, reduce_tmp);
        if (tid == 0) { sums[0] = sum0; sums[1] = sum1; }
        __syncthreads();
        float acc0[4] = {0.0f, 0.0f, 0.0f, 0.0f};
        float acc1[4] = {0.0f, 0.0f, 0.0f, 0.0f};
        const __nv_bfloat16* vbase = vc + kvh * a.max_seq * D;
        for (int ix = warp; ix < len; ix += W) {
            float wt0 = expf(scores[0][ix] - maxima[0]);
            float wt1 = expf(scores[1][ix] - maxima[1]);
            uint2 bits = __ldg(reinterpret_cast<const uint2*>(vbase + (begin + ix) * D + lane * 4));
            const __nv_bfloat16* x = reinterpret_cast<const __nv_bfloat16*>(&bits);
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                float xv = bf(x[j]);
                acc0[j] += wt0 * xv;
                acc1[j] += wt1 * xv;
            }
        }
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            acc[0][warp][j][lane] = acc0[j];
            acc[1][warp][j][lane] = acc1[j];
        }
        __syncthreads();
        if (tid < 2) {
            int head = qh0 + tid;
            partial_max[head * a.max_chunks + chunk] = maxima[tid];
            partial_sum[head * a.max_chunks + chunk] = sums[tid];
        }
        if (tid < 2 * 32) {
            int which = tid / 32;
            int lane_out = tid & 31;
            float val0 = 0.0f, val1 = 0.0f, val2 = 0.0f, val3 = 0.0f;
            #pragma unroll
            for (int ww = 0; ww < W; ++ww) {
                val0 += acc[which][ww][0][lane_out];
                val1 += acc[which][ww][1][lane_out];
                val2 += acc[which][ww][2][lane_out];
                val3 += acc[which][ww][3][lane_out];
            }
            int head = qh0 + which;
            int index = head * a.max_chunks + chunk;
            float* outp = partial_out + index * D + lane_out * 4;
            *reinterpret_cast<float4*>(outp) = make_float4(val0, val1, val2, val3);
        }
    }
    grid.sync();
}

__device__ void attention_finalize(
    cg::grid_group& grid, const KernelArgs& a, float* partial_max,
    float* partial_sum, float* partial_out, float* attn) {
    int chunks = (a.pos + 1 + CHUNK - 1) / CHUNK;
    int qh = int(blockIdx.x);
    if (qh < NQ) {
        __shared__ float tmp[W];
        __shared__ float scalar[2];
        int tid = threadIdx.x;
        float local_max = -FLT_MAX;
        for (int c = tid; c < chunks; c += T)
            local_max = fmaxf(local_max, partial_max[qh * a.max_chunks + c]);
        float mx = block_max(local_max, tmp);
        if (tid == 0) scalar[0] = mx;
        __syncthreads();
        mx = scalar[0];
        float local_sum = 0.0f;
        for (int c = tid; c < chunks; c += T) {
            int i = qh * a.max_chunks + c;
            local_sum += partial_sum[i] * expf(partial_max[i] - mx);
        }
        float total = block_sum(local_sum, tmp);
        if (tid == 0) scalar[1] = total;
        __syncthreads();
        total = scalar[1];
        for (int d = tid; d < D; d += T) {
            float val = 0.0f;
            for (int c = 0; c < chunks; ++c) {
                int i = qh * a.max_chunks + c;
                val += partial_out[i * D + d] * expf(partial_max[i] - mx);
            }
            attn[qh * D + d] = val / total;
        }
    }
    grid.sync();
}

__device__ void o_proj_norm_gate(
    cg::grid_group& grid, const KernelArgs& a, const LayerWeights& w,
    const float* attn, float* activations, float* residual, float* mlp) {
    __shared__ float normed[H];
    __shared__ float red[W];
    __shared__ float inv_rms;
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int hrows = (H + gridDim.x - 1) / gridDim.x;
    int begin = int(blockIdx.x) * hrows;
    int end = min(begin + hrows, H);
    for (int row = begin + warp; row < end; row += W) {
        float z = row_dot(w.o_proj + row * (NQ * D), attn, NQ * D);
        if (lane == 0) activations[row] = z + residual[row];
    }
    grid.sync();
    float local = 0.0f;
    for (int i = threadIdx.x; i < H; i += T) {
        float x = activations[i];
        if (blockIdx.x == 0) residual[i] = x;
        normed[i] = x;
        local += x * x;
    }
    float s = warp_sum(local);
    if (lane == 0) red[warp] = s;
    __syncthreads();
    if (warp == 0) {
        float z = lane < W ? red[lane] : 0.0f;
        z = warp_sum(z);
        if (lane == 0) inv_rms = rsqrtf(z / float(H) + 1.0e-6f);
    }
    __syncthreads();
    for (int i = threadIdx.x; i < H; i += T)
        normed[i] = normed[i] * inv_rms * bf(w.post_ln[i]);
    __syncthreads();

    int rows = (I + gridDim.x - 1) / gridDim.x;
    begin = int(blockIdx.x) * rows;
    end = min(begin + rows, I);
    for (int row = begin + warp; row < end; row += W) {
        float g = row_dot(w.gate_proj + row * H, normed, H);
        float u = row_dot(w.up_proj + row * H, normed, H);
        if (lane == 0) mlp[row] = (g / (1.0f + expf(-g))) * u;
    }
}

__device__ void down_project(
    cg::grid_group& grid, const KernelArgs& a, const LayerWeights& w,
    const float* mlp, const float* residual, __nv_bfloat16* dst) {
    grid.sync();
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int rows = (H + gridDim.x - 1) / gridDim.x;
    int begin = int(blockIdx.x) * rows;
    int end = min(begin + rows, H);
    for (int row = begin + warp; row < end; row += W) {
        float z = row_dot(w.down_proj + row * I, mlp, I);
        if (lane == 0) dst[row] = __float2bfloat16_rn(z + residual[row]);
    }
}

__global__ __launch_bounds__(T, 2)
void decode_kernel(KernelArgs a) {
    cg::grid_group grid = cg::this_grid();
    const int layer_stride = a.max_seq * D * NK;
    for (int layer = 0; layer < a.num_layers; ++layer) {
        const LayerWeights& w = a.weights[layer];
        const __nv_bfloat16* src;
        __nv_bfloat16* dst;
        if (layer == 0) src = a.input;
        else src = ((layer & 1) ? a.h0 : a.h1);
        if ((layer & 1) == 0) dst = a.h0;
        else dst = a.h1;
        qkv_project(grid, a, w, src, a.noise, a.q, a.k, a.v, a.residual,
                    layer == 0, layer != 0);
        qk_norm_rope(grid, a, w, a.q, a.k, a.v,
                     a.k_cache[layer], a.v_cache[layer]);
        attention_partials(grid, a, a.q, a.k_cache[layer], a.v_cache[layer],
                           a.partial_max, a.partial_sum, a.partial_out);
        attention_finalize(grid, a, a.partial_max, a.partial_sum,
                           a.partial_out, a.attn);
        o_proj_norm_gate(grid, a, w, a.attn, a.activations, a.residual, a.mlp);
        down_project(grid, a, w, a.mlp, a.residual, dst);
    }
}

static inline const __nv_bfloat16* bf16ptr(const at::Tensor& t) {
    return reinterpret_cast<const __nv_bfloat16*>(t.data_ptr<at::BFloat16>());
}
static inline __nv_bfloat16* bf16ptrw(const at::Tensor& t) {
    return reinterpret_cast<__nv_bfloat16*>(t.data_ptr<at::BFloat16>());
}
static inline float* fptr(const at::Tensor& t) { return t.data_ptr<float>(); }

void megaqwen_decode_cuda(
    at::Tensor input, at::Tensor noise, at::Tensor h0, at::Tensor h1,
    std::vector<at::Tensor> weights,
    std::vector<at::Tensor> k_caches, std::vector<at::Tensor> v_caches,
    at::Tensor rope_cos, at::Tensor rope_sin,
    at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor attn,
    at::Tensor activations, at::Tensor residual, at::Tensor mlp,
    at::Tensor partial_max, at::Tensor partial_sum, at::Tensor partial_out,
    int64_t num_layers, int64_t pos, int64_t max_seq, int64_t max_chunks) {
    TORCH_CHECK(input.is_cuda() && input.scalar_type() == at::kBFloat16, "input must be CUDA bfloat16");
    TORCH_CHECK(num_layers > 0 && num_layers <= MAX_LAYERS, "unsupported layer count");
    TORCH_CHECK(weights.size() == size_t(num_layers * 11), "weight list size mismatch");
    TORCH_CHECK(k_caches.size() == size_t(num_layers) && v_caches.size() == size_t(num_layers),
                "cache list size mismatch");
    c10::cuda::CUDAGuard guard(input.device());
    KernelArgs a{};
    a.input = bf16ptr(input); a.noise = bf16ptr(noise);
    a.h0 = bf16ptrw(h0); a.h1 = bf16ptrw(h1);
    for (int l = 0; l < num_layers; ++l) {
        int b = l * 11;
        a.weights[l] = {
            bf16ptr(weights[b + 0]), bf16ptr(weights[b + 1]), bf16ptr(weights[b + 2]),
            bf16ptr(weights[b + 3]), bf16ptr(weights[b + 4]), bf16ptr(weights[b + 5]),
            bf16ptr(weights[b + 6]), bf16ptr(weights[b + 7]), bf16ptr(weights[b + 8]),
            bf16ptr(weights[b + 9]), bf16ptr(weights[b + 10])};
        a.k_cache[l] = bf16ptrw(k_caches[l]);
        a.v_cache[l] = bf16ptrw(v_caches[l]);
    }
    a.rope_cos = fptr(rope_cos); a.rope_sin = fptr(rope_sin);
    a.q = fptr(q); a.k = fptr(k); a.v = fptr(v); a.attn = fptr(attn);
    a.activations = fptr(activations); a.residual = fptr(residual); a.mlp = fptr(mlp);
    a.partial_max = fptr(partial_max); a.partial_sum = fptr(partial_sum);
    a.partial_out = fptr(partial_out);
    a.num_layers = int(num_layers); a.pos = int(pos); a.max_seq = int(max_seq);
    a.max_chunks = int(max_chunks);
    int device = input.get_device();
    cudaDeviceProp prop;
    C10_CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
    int active = 0;
    C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &active, decode_kernel, T, 0));
    int blocks = prop.multiProcessorCount * active;
    void* args[] = {&a};
    cudaStream_t stream = at::cuda::getCurrentCUDAStream(device);
    C10_CUDA_CHECK(cudaLaunchCooperativeKernel(
        reinterpret_cast<void*>(decode_kernel), dim3(blocks), dim3(T), args, 0, stream));
}
"""


@lru_cache(maxsize=1)
def _extension():
    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
    from torch.utils.cpp_extension import load_inline

    return load_inline(
        name="megaqwen_decode_sm120_v1",
        cpp_sources=_CPP_SOURCE,
        cuda_sources=_CUDA_SOURCE,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3", "--use_fast_math"],
        with_cuda=True,
        verbose=False,
    )


@lru_cache(maxsize=8)
def _rope_tables(device_index: int, max_seq: int):
    device = torch.device("cuda", device_index)
    pos = torch.arange(max_seq, device=device, dtype=torch.float32)
    dim = torch.arange(HEAD_DIM // 2, device=device, dtype=torch.float32)
    inv = 1.0 / torch.pow(torch.tensor(10000.0, device=device), dim / (HEAD_DIM // 2))
    angles = pos[:, None] * inv[None, :]
    return angles.cos().contiguous(), angles.sin().contiguous()


def _scratch(device, max_seq):
    chunks = (max_seq + 255) // 256
    return (
        torch.empty((NUM_Q * HEAD_DIM,), device=device, dtype=torch.float32),
        torch.empty((NUM_KV * HEAD_DIM,), device=device, dtype=torch.float32),
        torch.empty((NUM_KV * HEAD_DIM,), device=device, dtype=torch.float32),
        torch.empty((NUM_Q * HEAD_DIM,), device=device, dtype=torch.float32),
        torch.empty((HIDDEN,), device=device, dtype=torch.float32),
        torch.empty((HIDDEN,), device=device, dtype=torch.float32),
        torch.empty((INTERMEDIATE,), device=device, dtype=torch.float32),
        torch.empty((NUM_Q * chunks,), device=device, dtype=torch.float32),
        torch.empty((NUM_Q * chunks,), device=device, dtype=torch.float32),
        torch.empty((NUM_Q * chunks * HEAD_DIM,), device=device, dtype=torch.float32),
    )


def _ordered_weights(model):
    names = (
        "input_ln", "q_proj", "k_proj", "v_proj", "q_norm", "k_norm",
        "o_proj", "post_ln", "gate_proj", "up_proj", "down_proj",
    )
    return [getattr(block, name) for block in model.blocks for name in names]


def _invoke(model, source, noise, h0, h1, k_caches, v_caches, pos, scratch, tables):
    cos, sin = tables
    max_seq = model.max_seq
    max_chunks = (max_seq + 255) // 256
    _extension().decode(
        source, noise, h0, h1, _ordered_weights(model), k_caches, v_caches,
        cos, sin, *scratch[:7], *scratch[7:],
        model.num_layers, pos, max_seq, max_chunks,
    )
    return h0 if model.num_layers & 1 else h1


def _noise_rows(count: int, seed: int) -> torch.Tensor:
    g = torch.Generator(device="cpu")
    g.manual_seed(seed)
    out = torch.empty((count, HIDDEN), dtype=torch.bfloat16, device="cpu")
    for i in range(count):
        out[i].copy_(torch.randn(HIDDEN, generator=g, dtype=torch.bfloat16))
    return out


def _initial_hidden(seed: int, device):
    g = torch.Generator(device="cpu")
    g.manual_seed(seed)
    return torch.randn(HIDDEN, generator=g, dtype=torch.bfloat16).to(device)


def _empty_caches(model, device):
    return (
        [torch.zeros(NUM_KV, model.max_seq, HEAD_DIM, device=device, dtype=torch.bfloat16)
         for _ in range(model.num_layers)],
        [torch.zeros(NUM_KV, model.max_seq, HEAD_DIM, device=device, dtype=torch.bfloat16)
         for _ in range(model.num_layers)],
    )


@torch.no_grad()
def prefill(model: Model, ctx_len: int, seed: int):
    device = next(model.parameters()).device
    model.eval()
    if ctx_len > model.max_seq:
        raise ValueError("ctx_len exceeds model.max_seq")
    h = _initial_hidden(seed, device)
    k_caches, v_caches = _empty_caches(model, device)
    noise = _noise_rows(ctx_len, seed + 1).to(device)
    scratch = _scratch(device, model.max_seq)
    h0 = torch.empty_like(h)
    h1 = torch.empty_like(h)
    tables = _rope_tables(device.index or 0, model.max_seq)
    for pos in range(ctx_len):
        h = _invoke(model, h, noise[pos], h0, h1, k_caches, v_caches, pos, scratch, tables)
    return h, k_caches, v_caches


@torch.no_grad()
def decode_steps(model, hidden, k_caches, v_caches, start_pos, n_steps, seed):
    if n_steps == 0:
        return hidden, k_caches, v_caches
    device = hidden.device
    model.eval()
    if start_pos + n_steps > model.max_seq:
        raise ValueError("decode range exceeds model.max_seq")
    noise = _noise_rows(n_steps, seed + 2).to(device)
    scratch = _scratch(device, model.max_seq)
    h0 = torch.empty_like(hidden)
    h1 = torch.empty_like(hidden)
    tables = _rope_tables(device.index or 0, model.max_seq)
    h = hidden
    for i in range(n_steps):
        h = _invoke(model, h, noise[i], h0, h1, k_caches, v_caches,
                    start_pos + i, scratch, tables)
    return h, k_caches, v_caches


@torch.no_grad()
def run(ctx_len, decode_steps, seed, model=None):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    max_seq = max(ctx_len + decode_steps, 512)
    if model is None:
        model = Model(NUM_LAYERS, max_seq)
    elif model.max_seq < ctx_len + decode_steps:
        raise ValueError("model.max_seq is too small for this run")
    model = model.to(device).eval()
    h, kc, vc = prefill(model, ctx_len, seed)
    h, kc, vc = globals()["decode_steps"](
        model, h, kc, vc, ctx_len, decode_steps, seed
    )
    return {"last_hidden": h}