MegaQwen Decode Grok 4.5

"""Fast multi-layer Qwen3-0.6B-geometry decode for RTX PRO 6000 (SM120).

Multi-kernel CUDA path (no cooperative grid.sync — better on long context and
fewer stalls than a full megakernel). Fused RMSNorm+QKV, QK-norm+RoPE+cache,
GQA attention with shared KV reads, fused O+residual+post-RMSNorm, gate/up
SiLU, down. Multi-step loop in C++. Matches reference.Model state_dict.
"""
from __future__ import annotations

import os
from typing import Optional

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_orig_normal_ = torch.nn.init.normal_


def _normal_compat(tensor, mean=0.0, std=1.0, generator=None):
    if (
        generator is not None
        and getattr(tensor, "is_cuda", False)
        and getattr(generator, "device", None) is not None
        and generator.device.type == "cpu"
    ):
        tmp = torch.empty(tensor.shape, dtype=tensor.dtype, device="cpu")
        _orig_normal_(tmp, mean=mean, std=std, generator=generator)
        with torch.no_grad():
            tensor.copy_(tmp)
        return tensor
    return _orig_normal_(tensor, mean=mean, std=std, generator=generator)


torch.nn.init.normal_ = _normal_compat  # type: ignore[assignment]

HIDDEN = 1024
INTERMEDIATE = 3072
NUM_Q = 16
NUM_KV = 8
HEAD_DIM = 128
NUM_LAYERS = 4
Q_SIZE = NUM_Q * HEAD_DIM
KV_SIZE = NUM_KV * HEAD_DIM

_mod = None


def _cuda_src() -> str:
    return r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <vector>
#include <algorithm>
#include <cmath>

constexpr int H = 1024;
constexpr int ISIZE = 3072;
constexpr int NQ = 16;
constexpr int NKV = 8;
constexpr int HD = 128;
constexpr int QS = 2048;
constexpr int KVS = 1024;
constexpr int WARP = 32;
constexpr float RMS_EPS = 1e-6f;
constexpr float ATTN_SCALE = 0.08838834764831843f;
constexpr int GQA = 2;
constexpr int MAX_LAYERS = 8;
constexpr int BLOCK = 256;
constexpr int NWARPS = BLOCK / WARP;

struct LayerW {
    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;
    const __nv_bfloat16* up;
    const __nv_bfloat16* down;
};

__device__ __forceinline__ float warp_sum(float v) {
#pragma unroll
    for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
    return v;
}
__device__ __forceinline__ float warp_max(float v) {
#pragma unroll
    for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_down_sync(0xffffffff, v, o));
    return v;
}
__device__ __forceinline__ float silu(float x) {
    return x * (1.f / (1.f + expf(-x)));
}

static inline int pick_blocks(int rows, int maxb) {
    int need = (rows + NWARPS - 1) / NWARPS;
    return std::max(1, std::min(need, maxb));
}

// -------------------- kernels --------------------

__global__ void __launch_bounds__(256, 2)
rms_qkv_kernel(
    const __nv_bfloat16* __restrict__ x,
    const __nv_bfloat16* __restrict__ w_norm,
    const __nv_bfloat16* __restrict__ Wq,
    const __nv_bfloat16* __restrict__ Wk,
    const __nv_bfloat16* __restrict__ Wv,
    float* __restrict__ q, float* __restrict__ k, float* __restrict__ v,
    float* __restrict__ residual
) {
    extern __shared__ float smem[];
    float* s_act = smem;
    float* s_red = smem + H;
    const int tid = threadIdx.x;
    const int warp = tid / WARP;
    const int lane = tid % WARP;

    float sq = 0.f;
    for (int i = tid; i < H; i += BLOCK) {
        float val = __bfloat162float(x[i]);
        s_act[i] = val;
        residual[i] = val;
        sq += val * val;
    }
    sq = warp_sum(sq);
    if (lane == 0) s_red[warp] = sq;
    __syncthreads();
    if (warp == 0) {
        float s = (lane < NWARPS) ? s_red[lane] : 0.f;
        s = warp_sum(s);
        if (lane == 0) s_red[0] = rsqrtf(s / float(H) + RMS_EPS);
    }
    __syncthreads();
    float rstd = s_red[0];
    for (int i = tid; i < H; i += BLOCK)
        s_act[i] *= rstd * __bfloat162float(__ldg(w_norm + i));
    __syncthreads();

    constexpr int TOTAL = QS + KVS + KVS;
    for (int m = blockIdx.x * NWARPS + warp; m < TOTAL; m += gridDim.x * NWARPS) {
        const __nv_bfloat16* row;
        float* outp;
        if (m < QS) { row = Wq + (size_t)m * H; outp = q + m; }
        else if (m < QS + KVS) { row = Wk + (size_t)(m - QS) * H; outp = k + (m - QS); }
        else { row = Wv + (size_t)(m - QS - KVS) * H; outp = v + (m - QS - KVS); }
        float sum = 0.f;
#pragma unroll 8
        for (int kk = lane * 4; kk < H; kk += WARP * 4) {
            uint2 wu = __ldg(reinterpret_cast<const uint2*>(row + kk));
            auto* wp = reinterpret_cast<__nv_bfloat16*>(&wu);
            sum += __bfloat162float(wp[0]) * s_act[kk]
                 + __bfloat162float(wp[1]) * s_act[kk+1]
                 + __bfloat162float(wp[2]) * s_act[kk+2]
                 + __bfloat162float(wp[3]) * s_act[kk+3];
        }
        sum = warp_sum(sum);
        if (lane == 0) *outp = sum;
    }
}

__global__ void qk_rope_cache_kernel(
    float* __restrict__ q, float* __restrict__ k, const float* __restrict__ v,
    const __nv_bfloat16* __restrict__ q_norm, const __nv_bfloat16* __restrict__ k_norm,
    __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache,
    int position, int max_seq
) {
    const int half = HD / 2;
    const int h = blockIdx.x;
    const int tid = threadIdx.x;

    auto process = [&](float* head, const __nv_bfloat16* nw, bool is_k, int kv_h) {
        __shared__ float sred[8];
        __shared__ float snorm[HD];
        float sq = 0.f;
        for (int i = tid; i < HD; i += blockDim.x) sq += head[i] * head[i];
        int lane = tid % WARP, warp = tid / WARP;
        float w = warp_sum(sq);
        if (lane == 0) sred[warp] = w;
        __syncthreads();
        if (warp == 0) {
            float s = (lane < (blockDim.x / WARP)) ? sred[lane] : 0.f;
            s = warp_sum(s);
            if (lane == 0) sred[0] = rsqrtf(s / float(HD) + RMS_EPS);
        }
        __syncthreads();
        float rstd = sred[0];
        for (int i = tid; i < HD; i += blockDim.x)
            snorm[i] = head[i] * rstd * __bfloat162float(__ldg(nw + i));
        __syncthreads();
        for (int i = tid; i < half; i += blockDim.x) {
            float inv = expf(-logf(10000.f) * (float)i / (float)half);
            float ang = (float)position * inv;
            float c = cosf(ang), s = sinf(ang);
            float x1 = snorm[i], x2 = snorm[i + half];
            float y1 = x1 * c - x2 * s;
            float y2 = x1 * s + x2 * c;
            head[i] = y1; head[i + half] = y2;
            if (is_k) {
                __nv_bfloat16* kc = k_cache + ((size_t)kv_h * max_seq + position) * HD;
                kc[i] = __float2bfloat16(y1);
                kc[i + half] = __float2bfloat16(y2);
            }
        }
        if (is_k) {
            __syncthreads();
            __nv_bfloat16* vc = v_cache + ((size_t)kv_h * max_seq + position) * HD;
            const float* vh = v + kv_h * HD;
            for (int i = tid; i < HD; i += blockDim.x)
                vc[i] = __float2bfloat16(vh[i]);
        }
    };
    if (h < NQ) process(q + h * HD, q_norm, false, 0);
    if (h < NKV) process(k + h * HD, k_norm, true, h);
}

// Multi-chunk GQA partial path (shared KV read for 2 Q heads per block)
__global__ void __launch_bounds__(256, 2)
attn_gqa_partial_kernel(
    const float* __restrict__ q,
    const __nv_bfloat16* __restrict__ k_cache,
    const __nv_bfloat16* __restrict__ v_cache,
    float* __restrict__ partial_m,
    float* __restrict__ partial_l,
    float* __restrict__ partial_o,
    int cache_len, int max_seq, int num_chunks, int chunk_stride
) {
    const int kv_h = blockIdx.x;
    const int chunk = blockIdx.y;
    const int tid = threadIdx.x;
    const int warp = tid / WARP;
    const int lane = tid % WARP;

    int chunk_sz = (cache_len + num_chunks - 1) / num_chunks;
    int p0 = chunk * chunk_sz;
    int p1 = min(p0 + chunk_sz, cache_len);
    const int qh0 = kv_h * GQA;
    const int qh1 = qh0 + 1;

    __shared__ float sq0[HD], sq1[HD];
    for (int d = tid; d < HD; d += BLOCK) {
        sq0[d] = q[qh0 * HD + d];
        sq1[d] = q[qh1 * HD + d];
    }
    __syncthreads();

    auto empty = [&](int qh) {
        if (tid == 0) {
            partial_m[qh * chunk_stride + chunk] = -INFINITY;
            partial_l[qh * chunk_stride + chunk] = 0.f;
        }
        for (int d = tid; d < HD; d += BLOCK)
            partial_o[(qh * chunk_stride + chunk) * HD + d] = 0.f;
    };
    if (p0 >= cache_len) { empty(qh0); empty(qh1); return; }

    float m0 = -INFINITY, l0 = 0.f, m1 = -INFINITY, l1 = 0.f;
    float o0[4] = {0,0,0,0}, o1[4] = {0,0,0,0};

    for (int pos = p0 + warp; pos < p1; pos += NWARPS) {
        const __nv_bfloat16* kp = k_cache + ((size_t)kv_h * max_seq + pos) * HD;
        const __nv_bfloat16* vp = v_cache + ((size_t)kv_h * max_seq + pos) * HD;
        float s0 = 0.f, s1 = 0.f;
#pragma unroll
        for (int d = lane * 2; d < HD; d += WARP * 2) {
            uint32_t ku = __ldg(reinterpret_cast<const uint32_t*>(kp + d));
            auto* kb = reinterpret_cast<__nv_bfloat16*>(&ku);
            float kf0 = __bfloat162float(kb[0]);
            float kf1 = __bfloat162float(kb[1]);
            s0 += sq0[d] * kf0 + sq0[d+1] * kf1;
            s1 += sq1[d] * kf0 + sq1[d+1] * kf1;
        }
        s0 = warp_sum(s0) * ATTN_SCALE;
        s1 = warp_sum(s1) * ATTN_SCALE;
        s0 = __shfl_sync(0xffffffff, s0, 0);
        s1 = __shfl_sync(0xffffffff, s1, 0);
        float vreg[4];
#pragma unroll
        for (int j = 0, d = lane; d < HD; d += WARP, ++j)
            vreg[j] = __bfloat162float(__ldg(vp + d));
        {
            float mnew = fmaxf(m0, s0);
            float ed = expf(m0 - mnew);
            float e = expf(s0 - mnew);
            l0 = l0 * ed + e;
#pragma unroll
            for (int j = 0; j < 4; ++j) o0[j] = o0[j] * ed + e * vreg[j];
            m0 = mnew;
        }
        {
            float mnew = fmaxf(m1, s1);
            float ed = expf(m1 - mnew);
            float e = expf(s1 - mnew);
            l1 = l1 * ed + e;
#pragma unroll
            for (int j = 0; j < 4; ++j) o1[j] = o1[j] * ed + e * vreg[j];
            m1 = mnew;
        }
    }

    __shared__ float sm0[8], sl0[8], sm1[8], sl1[8];
    __shared__ float so0[8][HD], so1[8][HD];
    if (lane == 0) { sm0[warp]=m0; sl0[warp]=l0; sm1[warp]=m1; sl1[warp]=l1; }
#pragma unroll
    for (int j = 0, d = lane; d < HD; d += WARP, ++j) {
        so0[warp][d] = o0[j]; so1[warp][d] = o1[j];
    }
    __syncthreads();

    auto combine = [&](int qh, float* sm, float* sl, float so[][HD]) {
        if (warp != 0) return;
        float gm = -INFINITY;
        for (int w = 0; w < NWARPS; ++w)
            if (sm[w] > -1e30f) gm = fmaxf(gm, sm[w]);
        float gl = 0.f; float go[4] = {0,0,0,0};
        for (int w = 0; w < NWARPS; ++w) {
            if (!(sm[w] > -1e30f)) continue;
            float sc = expf(sm[w] - gm);
            gl += sl[w] * sc;
#pragma unroll
            for (int j = 0, d = lane; d < HD; d += WARP, ++j)
                go[j] += so[w][d] * sc;
        }
        if (lane == 0) {
            partial_m[qh * chunk_stride + chunk] = gm;
            partial_l[qh * chunk_stride + chunk] = gl;
        }
#pragma unroll
        for (int j = 0, d = lane; d < HD; d += WARP, ++j)
            partial_o[(qh * chunk_stride + chunk) * HD + d] = go[j];
    };
    combine(qh0, sm0, sl0, so0);
    __syncthreads();
    combine(qh1, sm1, sl1, so1);
}

__global__ void attn_reduce_kernel(
    const float* __restrict__ partial_m,
    const float* __restrict__ partial_l,
    const float* __restrict__ partial_o,
    float* __restrict__ attn_out,
    int num_chunks, int chunk_stride
) {
    const int qh = blockIdx.x;
    const int tid = threadIdx.x;
    float gm = -INFINITY;
    for (int c = tid; c < num_chunks; c += blockDim.x)
        gm = fmaxf(gm, partial_m[qh * chunk_stride + c]);
    __shared__ float sred[8];
    int lane = tid % WARP, warp = tid / WARP;
    float wm = warp_max(gm);
    if (lane == 0) sred[warp] = wm;
    __syncthreads();
    if (warp == 0) {
        float v = (lane < blockDim.x / WARP) ? sred[lane] : -INFINITY;
        v = warp_max(v);
        if (lane == 0) sred[0] = v;
    }
    __syncthreads();
    gm = sred[0];
    for (int d = tid; d < HD; d += blockDim.x) {
        float gl = 0.f, go = 0.f;
        for (int c = 0; c < num_chunks; ++c) {
            float m = partial_m[qh * chunk_stride + c];
            if (!(m > -1e30f)) continue;
            float sc = expf(m - gm);
            gl += partial_l[qh * chunk_stride + c] * sc;
            go += partial_o[(qh * chunk_stride + c) * HD + d] * sc;
        }
        attn_out[qh * HD + d] = go / fmaxf(gl, 1e-20f);
    }
}

// Fused: O proj + residual → resid2, then RMSNorm(resid2) → normed
// Implemented as two phases in one kernel with block-local then grid-wide via two launches
// is messy. Keep separate but fuse O+write and do rms as single-block (tiny).

__global__ void __launch_bounds__(256, 2)
o_proj_kernel(
    const float* __restrict__ attn, const __nv_bfloat16* __restrict__ Wo,
    const float* __restrict__ residual, float* __restrict__ out
) {
    extern __shared__ float s_attn[];
    const int tid = threadIdx.x, warp = tid / WARP, lane = tid % WARP;
    for (int i = tid; i < QS; i += BLOCK) s_attn[i] = attn[i];
    __syncthreads();
    for (int m = blockIdx.x * NWARPS + warp; m < H; m += gridDim.x * NWARPS) {
        const __nv_bfloat16* row = Wo + (size_t)m * QS;
        float sum = 0.f;
#pragma unroll 8
        for (int k = lane * 4; k < QS; k += WARP * 4) {
            uint2 wu = __ldg(reinterpret_cast<const uint2*>(row + k));
            auto* wp = reinterpret_cast<__nv_bfloat16*>(&wu);
            sum += __bfloat162float(wp[0])*s_attn[k]+__bfloat162float(wp[1])*s_attn[k+1]
                 +__bfloat162float(wp[2])*s_attn[k+2]+__bfloat162float(wp[3])*s_attn[k+3];
        }
        sum = warp_sum(sum);
        if (lane == 0) out[m] = sum + residual[m];
    }
}

__global__ void rmsnorm_f32_kernel(
    const float* __restrict__ x, const __nv_bfloat16* __restrict__ w, float* __restrict__ y
) {
    __shared__ float sred[8];
    int tid = threadIdx.x, lane = tid % WARP, warp = tid / WARP;
    float sq = 0.f;
    for (int i = tid; i < H; i += BLOCK) { float v = x[i]; sq += v * v; }
    sq = warp_sum(sq);
    if (lane == 0) sred[warp] = sq;
    __syncthreads();
    if (warp == 0) {
        float s = (lane < NWARPS) ? sred[lane] : 0.f;
        s = warp_sum(s);
        if (lane == 0) sred[0] = rsqrtf(s / float(H) + RMS_EPS);
    }
    __syncthreads();
    float rstd = sred[0];
    for (int i = tid; i < H; i += BLOCK)
        y[i] = x[i] * rstd * __bfloat162float(__ldg(w + i));
}

__global__ void __launch_bounds__(256, 2)
gate_up_kernel(
    const float* __restrict__ x,
    const __nv_bfloat16* __restrict__ Wg, const __nv_bfloat16* __restrict__ Wu,
    float* __restrict__ out
) {
    extern __shared__ float s_act[];
    const int tid = threadIdx.x, warp = tid / WARP, lane = tid % WARP;
    for (int i = tid; i < H; i += BLOCK) s_act[i] = x[i];
    __syncthreads();
    for (int m = blockIdx.x * NWARPS + warp; m < ISIZE; m += gridDim.x * NWARPS) {
        const __nv_bfloat16* rg = Wg + (size_t)m * H;
        const __nv_bfloat16* ru = Wu + (size_t)m * H;
        float sg = 0.f, su = 0.f;
#pragma unroll 8
        for (int k = lane * 4; k < H; k += WARP * 4) {
            uint2 gu = __ldg(reinterpret_cast<const uint2*>(rg + k));
            uint2 uu = __ldg(reinterpret_cast<const uint2*>(ru + k));
            auto* gp = reinterpret_cast<__nv_bfloat16*>(&gu);
            auto* up = reinterpret_cast<__nv_bfloat16*>(&uu);
            float a0=s_act[k],a1=s_act[k+1],a2=s_act[k+2],a3=s_act[k+3];
            sg += __bfloat162float(gp[0])*a0+__bfloat162float(gp[1])*a1
                +__bfloat162float(gp[2])*a2+__bfloat162float(gp[3])*a3;
            su += __bfloat162float(up[0])*a0+__bfloat162float(up[1])*a1
                +__bfloat162float(up[2])*a2+__bfloat162float(up[3])*a3;
        }
        sg = warp_sum(sg); su = warp_sum(su);
        if (lane == 0) out[m] = silu(sg) * su;
    }
}

__global__ void __launch_bounds__(256, 2)
down_residual_kernel(
    const float* __restrict__ mid, const __nv_bfloat16* __restrict__ Wd,
    const float* __restrict__ residual, __nv_bfloat16* __restrict__ out
) {
    extern __shared__ float s_mid[];
    const int tid = threadIdx.x, warp = tid / WARP, lane = tid % WARP;
    for (int i = tid; i < ISIZE; i += BLOCK) s_mid[i] = mid[i];
    __syncthreads();
    for (int m = blockIdx.x * NWARPS + warp; m < H; m += gridDim.x * NWARPS) {
        const __nv_bfloat16* row = Wd + (size_t)m * ISIZE;
        float sum = 0.f;
#pragma unroll 8
        for (int k = lane * 4; k < ISIZE; k += WARP * 4) {
            uint2 wu = __ldg(reinterpret_cast<const uint2*>(row + k));
            auto* wp = reinterpret_cast<__nv_bfloat16*>(&wu);
            sum += __bfloat162float(wp[0])*s_mid[k]+__bfloat162float(wp[1])*s_mid[k+1]
                 +__bfloat162float(wp[2])*s_mid[k+2]+__bfloat162float(wp[3])*s_mid[k+3];
        }
        sum = warp_sum(sum);
        if (lane == 0) out[m] = __float2bfloat16(sum + residual[m]);
    }
}

__global__ void mix_hidden_kernel(
    const __nv_bfloat16* __restrict__ noise,
    const __nv_bfloat16* __restrict__ h,
    __nv_bfloat16* __restrict__ out
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < H) {
        float a = __bfloat162float(noise[i]);
        float b = __bfloat162float(h[i]);
        out[i] = __float2bfloat16(0.5f * a + 0.5f * b);
    }
}

// -------------------- host --------------------

static int choose_chunks(int cache_len, int chunk_stride, int max_blocks) {
    // Always spread work across SMs: NKV * chunks ≈ max_blocks when possible.
    // Each chunk should still have enough positions to amortize launch (~32+).
    int sm_chunks = std::max(1, max_blocks / NKV);  // e.g. 188/8 ≈ 23
    int by_len = std::max(1, (cache_len + 63) / 64); // ≥1 chunk per 64 positions
    int n = std::min(chunk_stride, std::min(sm_chunks, by_len));
    // Prefer multi-chunk path whenever we can put >1 chunk of real work
    if (cache_len >= 128) n = std::max(n, std::min(chunk_stride, 4));
    if (cache_len >= 512) n = std::max(n, std::min(chunk_stride, 8));
    if (cache_len >= 2048) n = std::max(n, std::min(chunk_stride, 16));
    if (cache_len >= 8192) n = std::max(n, std::min(chunk_stride, sm_chunks));
    return std::max(1, std::min(n, chunk_stride));
}

static void launch_layer(
    const __nv_bfloat16* x, const LayerW& w,
    __nv_bfloat16* k_cache, __nv_bfloat16* v_cache,
    __nv_bfloat16* y,
    float* residual, float* g_q, float* g_k, float* g_v, float* g_attn,
    float* g_normed, float* g_mid, float* g_resid2,
    float* partial_m, float* partial_l, float* partial_o,
    int position, int cache_len, int max_seq,
    int num_chunks, int chunk_stride, int max_blocks,
    cudaStream_t stream
) {
    size_t sh = (H + NWARPS) * sizeof(float);
    size_t sq = QS * sizeof(float);
    size_t si = ISIZE * sizeof(float);

    rms_qkv_kernel<<<pick_blocks(QS+KVS+KVS, max_blocks), BLOCK, sh, stream>>>(
        x, w.input_ln, w.q_proj, w.k_proj, w.v_proj, g_q, g_k, g_v, residual);
    qk_rope_cache_kernel<<<NQ, 128, 0, stream>>>(
        g_q, g_k, g_v, w.q_norm, w.k_norm, k_cache, v_cache, position, max_seq);

    // Always use multi-block GQA path so long-seq and medium-seq both fill SMs.
    // (Single-block-per-KV-head full path under-utilizes a 188-SM GPU.)
    {
        dim3 grid(NKV, num_chunks);
        attn_gqa_partial_kernel<<<grid, BLOCK, 0, stream>>>(
            g_q, k_cache, v_cache, partial_m, partial_l, partial_o,
            cache_len, max_seq, num_chunks, chunk_stride);
        attn_reduce_kernel<<<NQ, 128, 0, stream>>>(
            partial_m, partial_l, partial_o, g_attn, num_chunks, chunk_stride);
    }

    o_proj_kernel<<<pick_blocks(H, max_blocks), BLOCK, sq, stream>>>(
        g_attn, w.o_proj, residual, g_resid2);
    rmsnorm_f32_kernel<<<1, BLOCK, 0, stream>>>(g_resid2, w.post_ln, g_normed);
    // residual for MLP = g_resid2 (post-attn). Keep g_resid2; don't clobber.
    gate_up_kernel<<<pick_blocks(ISIZE, max_blocks), BLOCK, H*sizeof(float), stream>>>(
        g_normed, w.gate, w.up, g_mid);
    down_residual_kernel<<<pick_blocks(H, max_blocks), BLOCK, si, stream>>>(
        g_mid, w.down, g_resid2, y);
}

static LayerW make_layer(const std::vector<torch::Tensor>& wf, int layer) {
    auto p = [&](int i) {
        return reinterpret_cast<const __nv_bfloat16*>(wf[layer * 11 + i].data_ptr());
    };
    return LayerW{p(0),p(1),p(2),p(3),p(4),p(5),p(6),p(7),p(8),p(9),p(10)};
}

// One decode step (mix + layers + writeback). position/cache_len/num_chunks fixed.
static void one_step(
    const __nv_bfloat16* noise_row,
    __nv_bfloat16* h_ptr,
    __nv_bfloat16* ba, __nv_bfloat16* bb,
    LayerW* layers, int num_layers,
    __nv_bfloat16** kptrs, __nv_bfloat16** vptrs,
    float* residual, float* g_q, float* g_k, float* g_v, float* g_attn,
    float* g_normed, float* g_mid, float* g_resid2,
    float* partial_m, float* partial_l, float* partial_o,
    int pos, int cache_len, int max_seq, int num_chunks, int chunk_stride,
    int max_blocks, cudaStream_t stream
) {
    mix_hidden_kernel<<<(H + 255) / 256, 256, 0, stream>>>(noise_row, h_ptr, ba);
    const __nv_bfloat16* cur = ba;
    __nv_bfloat16* dst = bb;
    for (int layer = 0; layer < num_layers; ++layer) {
        dst = (cur == ba) ? bb : ba;
        launch_layer(
            cur, layers[layer], kptrs[layer], vptrs[layer], dst,
            residual, g_q, g_k, g_v, g_attn, g_normed, g_mid, g_resid2,
            partial_m, partial_l, partial_o,
            pos, cache_len, max_seq, num_chunks, chunk_stride, max_blocks, stream
        );
        cur = dst;
    }
    if (cur != h_ptr) {
        cudaMemcpyAsync(h_ptr, cur, H * sizeof(__nv_bfloat16),
                        cudaMemcpyDeviceToDevice, stream);
    }
}

void run_steps_cuda(
    torch::Tensor hidden,
    torch::Tensor noise,
    const std::vector<torch::Tensor>& weights_flat,
    std::vector<torch::Tensor> k_caches,
    std::vector<torch::Tensor> v_caches,
    torch::Tensor buf_a,
    torch::Tensor buf_b,
    torch::Tensor residual,
    torch::Tensor g_q, torch::Tensor g_k, torch::Tensor g_v,
    torch::Tensor g_attn, torch::Tensor g_normed, torch::Tensor g_mid,
    torch::Tensor g_resid2,
    torch::Tensor partial_m, torch::Tensor partial_l, torch::Tensor partial_o,
    int start_pos, int n_steps, int max_seq, int num_layers, int max_blocks
) {
    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
    int chunk_stride = (int)partial_m.size(1);

    LayerW layers[MAX_LAYERS];
    for (int i = 0; i < num_layers; ++i) layers[i] = make_layer(weights_flat, i);

    __nv_bfloat16* kptrs[MAX_LAYERS];
    __nv_bfloat16* vptrs[MAX_LAYERS];
    for (int i = 0; i < num_layers; ++i) {
        kptrs[i] = reinterpret_cast<__nv_bfloat16*>(k_caches[i].data_ptr());
        vptrs[i] = reinterpret_cast<__nv_bfloat16*>(v_caches[i].data_ptr());
    }

    auto* h_ptr = reinterpret_cast<__nv_bfloat16*>(hidden.data_ptr());
    auto* noise_ptr = reinterpret_cast<__nv_bfloat16*>(noise.data_ptr());
    auto* ba = reinterpret_cast<__nv_bfloat16*>(buf_a.data_ptr());
    auto* bb = reinterpret_cast<__nv_bfloat16*>(buf_b.data_ptr());
    float* p_res = residual.data_ptr<float>();
    float* p_q = g_q.data_ptr<float>();
    float* p_k = g_k.data_ptr<float>();
    float* p_v = g_v.data_ptr<float>();
    float* p_attn = g_attn.data_ptr<float>();
    float* p_norm = g_normed.data_ptr<float>();
    float* p_mid = g_mid.data_ptr<float>();
    float* p_r2 = g_resid2.data_ptr<float>();
    float* p_pm = partial_m.data_ptr<float>();
    float* p_pl = partial_l.data_ptr<float>();
    float* p_po = partial_o.data_ptr<float>();

    // Group consecutive steps that share the same num_chunks; within a group
    // cache_len grows so we re-issue kernels (position is a kernel arg) but
    // still avoid Python overhead. CUDA graphs are awkward here because
    // position/cache_len change every step — rely on low-overhead launches.
    for (int step = 0; step < n_steps; ++step) {
        int pos = start_pos + step;
        int cache_len = pos + 1;
        int num_chunks = choose_chunks(cache_len, chunk_stride, max_blocks);
        one_step(
            noise_ptr + (size_t)step * H, h_ptr, ba, bb,
            layers, num_layers, kptrs, vptrs,
            p_res, p_q, p_k, p_v, p_attn, p_norm, p_mid, p_r2,
            p_pm, p_pl, p_po,
            pos, cache_len, max_seq, num_chunks, chunk_stride, max_blocks, stream
        );
    }
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run_steps_cuda", &run_steps_cuda);
}
'''


def _get_mod():
    global _mod
    if _mod is not None:
        return _mod
    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
    extra = [
        "-O3",
        "--use_fast_math",
        "-lineinfo",
        "-U__CUDA_NO_HALF_OPERATORS__",
        "-U__CUDA_NO_HALF_CONVERSIONS__",
        "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
        "-U__CUDA_NO_BFLOAT16_OPERATORS__",
        "--expt-relaxed-constexpr",
        "-std=c++17",
    ]
    try:
        major, minor = torch.cuda.get_device_capability(0)
        extra.append(f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}")
    except Exception:
        pass
    _mod = load_inline(
        name="megaqwen_decode_v5",
        cpp_sources=[],
        cuda_sources=[_cuda_src()],
        extra_cuda_cflags=extra,
        verbose=False,
    )
    return _mod


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        H, I, D = HIDDEN, INTERMEDIATE, HEAD_DIM
        self.input_ln = nn.Parameter(torch.ones(H, dtype=torch.bfloat16))
        self.q_proj = nn.Parameter(torch.empty(NUM_Q * D, H, dtype=torch.bfloat16))
        self.k_proj = nn.Parameter(torch.empty(NUM_KV * D, H, dtype=torch.bfloat16))
        self.v_proj = nn.Parameter(torch.empty(NUM_KV * D, H, dtype=torch.bfloat16))
        self.q_norm = nn.Parameter(torch.ones(D, dtype=torch.bfloat16))
        self.k_norm = nn.Parameter(torch.ones(D, dtype=torch.bfloat16))
        self.o_proj = nn.Parameter(torch.empty(H, NUM_Q * D, dtype=torch.bfloat16))
        self.post_ln = nn.Parameter(torch.ones(H, dtype=torch.bfloat16))
        self.gate_proj = nn.Parameter(torch.empty(I, H, dtype=torch.bfloat16))
        self.up_proj = nn.Parameter(torch.empty(I, H, dtype=torch.bfloat16))
        self.down_proj = nn.Parameter(torch.empty(H, I, dtype=torch.bfloat16))
        for p in self.parameters():
            if p is self.input_ln or p is self.post_ln or p is self.q_norm or p is self.k_norm:
                continue
            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)])
        self._ws = None
        self._max_blocks = None

    def _ensure_ws(self, device):
        if self._ws is not None and self._ws["residual"].device == device:
            return self._ws
        props = torch.cuda.get_device_properties(device)
        self._max_blocks = int(props.multi_processor_count)
        max_chunks = max(64, min(256, max(32, self._max_blocks // NUM_KV)))
        self._ws = {
            "residual": torch.empty(HIDDEN, device=device, dtype=torch.float32),
            "g_q": torch.empty(Q_SIZE, device=device, dtype=torch.float32),
            "g_k": torch.empty(KV_SIZE, device=device, dtype=torch.float32),
            "g_v": torch.empty(KV_SIZE, device=device, dtype=torch.float32),
            "g_attn": torch.empty(Q_SIZE, device=device, dtype=torch.float32),
            "g_normed": torch.empty(HIDDEN, device=device, dtype=torch.float32),
            "g_mid": torch.empty(INTERMEDIATE, device=device, dtype=torch.float32),
            "g_resid2": torch.empty(HIDDEN, device=device, dtype=torch.float32),
            "partial_m": torch.empty(NUM_Q, max_chunks, device=device, dtype=torch.float32),
            "partial_l": torch.empty(NUM_Q, max_chunks, device=device, dtype=torch.float32),
            "partial_o": torch.empty(NUM_Q, max_chunks, HEAD_DIM, device=device, dtype=torch.float32),
            "buf_a": torch.empty(HIDDEN, device=device, dtype=torch.bfloat16),
            "buf_b": torch.empty(HIDDEN, device=device, dtype=torch.bfloat16),
        }
        return self._ws

    def weight_list(self):
        out = []
        for b in self.blocks:
            out.extend(
                [
                    b.input_ln,
                    b.q_proj,
                    b.k_proj,
                    b.v_proj,
                    b.q_norm,
                    b.k_norm,
                    b.o_proj,
                    b.post_ln,
                    b.gate_proj,
                    b.up_proj,
                    b.down_proj,
                ]
            )
        return out


def empty_caches(num_layers: int, max_seq: int, device, dtype=torch.bfloat16):
    k = [
        torch.zeros(NUM_KV, max_seq, HEAD_DIM, device=device, dtype=dtype)
        for _ in range(num_layers)
    ]
    v = [
        torch.zeros(NUM_KV, max_seq, HEAD_DIM, device=device, dtype=dtype)
        for _ in range(num_layers)
    ]
    return k, v


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


def _run_steps(model: Model, hidden: torch.Tensor, k_caches, v_caches,
               start_pos: int, n_steps: int, noise: torch.Tensor) -> torch.Tensor:
    mod = _get_mod()
    device = hidden.device
    ws = model._ensure_ws(device)
    h_work = torch.empty(HIDDEN, device=device, dtype=torch.bfloat16)
    h_work.copy_(hidden)
    mod.run_steps_cuda(
        h_work,
        noise.contiguous(),
        model.weight_list(),
        k_caches,
        v_caches,
        ws["buf_a"],
        ws["buf_b"],
        ws["residual"],
        ws["g_q"],
        ws["g_k"],
        ws["g_v"],
        ws["g_attn"],
        ws["g_normed"],
        ws["g_mid"],
        ws["g_resid2"],
        ws["partial_m"],
        ws["partial_l"],
        ws["partial_o"],
        int(start_pos),
        int(n_steps),
        int(model.max_seq),
        int(model.num_layers),
        int(model._max_blocks),
    )
    return h_work


@torch.no_grad()
def prefill(model: Model, ctx_len: int, seed: int, device=None):
    device = device or next(model.parameters()).device
    model = model.to(device).eval()
    assert ctx_len <= model.max_seq
    _get_mod()
    h = _seeded_hidden(seed, device)
    k_caches, v_caches = empty_caches(model.num_layers, model.max_seq, device)
    g = torch.Generator(device="cpu")
    g.manual_seed(seed + 1)
    noise = torch.randn(ctx_len, HIDDEN, generator=g, dtype=torch.bfloat16, device="cpu").to(device)
    h = _run_steps(model, h, k_caches, v_caches, 0, ctx_len, noise)
    return h, k_caches, v_caches


@torch.no_grad()
def decode_steps(
    model: Model,
    hidden: torch.Tensor,
    k_caches,
    v_caches,
    start_pos: int,
    n_steps: int,
    seed: int,
):
    g = torch.Generator(device="cpu")
    g.manual_seed(seed + 2)
    noise = torch.randn(n_steps, HIDDEN, generator=g, dtype=torch.bfloat16, device="cpu").to(
        hidden.device
    )
    h = _run_steps(model, hidden, k_caches, v_caches, start_pos, n_steps, noise)
    return h, k_caches, v_caches


def run(
    ctx_len: int,
    n_decode: int,
    seed: int,
    model: Optional[Model] = None,
    max_seq: Optional[int] = None,
) -> dict:
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    max_seq = max_seq or max(ctx_len + n_decode, 512)
    if model is None:
        model = Model(NUM_LAYERS, max_seq)
    model = model.to(device).eval()
    h, k_caches, v_caches = prefill(model, ctx_len, seed, device=device)
    h, k_caches, v_caches = decode_steps(
        model, h, k_caches, v_caches, start_pos=ctx_len, n_steps=n_decode, seed=seed
    )
    return {
        "last_hidden": h.detach(),
        "ctx_len": ctx_len,
        "decode_steps": n_decode,
    }