"""Single cooperative-grid W4A16 hybrid decode, with absorbed latent attention."""

from dataclasses import dataclass, field
import os

import ninja

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


@dataclass(frozen=True)
class Config:
    hidden: int = 2304
    kda_heads: int = 32
    kda_head_dim: int = 128
    short_conv: int = 4
    mla_heads: int = 32
    kv_lora: int = 512
    qk_nope: int = 128
    qk_rope: int = 64
    v_head: int = 128
    rope_theta: float = 10000.0
    n_experts: int = 64
    n_active: int = 8
    n_shared: int = 1
    moe_inter: int = 1024
    routed_scaling: float = 2.446
    group: int = 128
    pattern: tuple = ("K", "K", "K", "M")
    dtype: torch.dtype = field(default=torch.bfloat16)


def build_config(shape):
    return Config(n_experts=int(shape.get("n_experts", 64)))


class QuantLinear(nn.Module):
    def __init__(self, in_f, out_f, group=128, experts=None):
        super().__init__()
        leading = () if experts is None else (experts,)
        self.register_buffer("w_q", torch.empty(*leading, in_f // 2, out_f, dtype=torch.uint8))
        self.register_buffer("scales", torch.empty(*leading, in_f // group, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.empty(*leading, in_f // group, out_f, dtype=torch.bfloat16))


class KDA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        for name in ("q_proj", "k_proj", "v_proj", "g_proj"):
            setattr(self, name, QuantLinear(cfg.hidden, cfg.kda_heads * cfg.kda_head_dim))
        self.beta_proj = nn.Linear(cfg.hidden, cfg.kda_heads, bias=False, dtype=cfg.dtype)
        self.conv_w = nn.Parameter(torch.empty(3, 4096, 4, dtype=cfg.dtype))
        self.o_proj = QuantLinear(4096, cfg.hidden)


class MLA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.q_proj = QuantLinear(cfg.hidden, 6144)
        self.kv_a = QuantLinear(cfg.hidden, 576)
        self.kv_b = QuantLinear(512, 8192)
        self.o_proj = QuantLinear(4096, cfg.hidden)


class MoE(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.router = nn.Linear(cfg.hidden, cfg.n_experts, bias=False, dtype=cfg.dtype)
        for prefix, count in (("", cfg.n_experts), ("s_", cfg.n_shared)):
            setattr(self, prefix + "gate", QuantLinear(2304, 1024, experts=count))
            setattr(self, prefix + "up", QuantLinear(2304, 1024, experts=count))
            setattr(self, prefix + "down", QuantLinear(1024, 2304, experts=count))


class Block(nn.Module):
    def __init__(self, cfg, kind):
        super().__init__()
        self.attn_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.moe_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.attn = KDA(cfg) if kind == "K" else MLA(cfg)
        self.moe = MoE(cfg)


_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
#include <mma.h>
#include <vector>

using BFloat = __nv_bfloat16;
namespace cg = cooperative_groups;
namespace wm = nvcuda::wmma;

struct Quant {
    const unsigned char* weight;
    const BFloat* scales;
    const BFloat* zeros;
};

struct Layer {
    const BFloat* attn_norm;
    const BFloat* moe_norm;
    const BFloat* beta;
    const BFloat* conv;
    const BFloat* router;
    Quant attn[5];
    Quant moe[6];
};

struct State {
    const float* old_state;
    float* new_state;
    const BFloat* old_conv[3];
    BFloat* new_conv[3];
};

struct Arguments {
    Layer layers[4];
    State states[3];
    const BFloat* hidden;
    BFloat* output;
    float* workspace;
    const BFloat* old_latent;
    const BFloat* old_rope;
    BFloat* latent;
    BFloat* rope;
    int context;
    int experts;
    float routing_scale;
};

constexpr int X_OFFSET = 0;
constexpr int H_OFFSET = 4096;
constexpr int PART_OFFSET = 8192;
constexpr int AUX_OFFSET = PART_OFFSET + 200704;
constexpr int ROUTE_OFFSET = AUX_OFFSET + 16384;
constexpr int QLAT_OFFSET = ROUTE_OFFSET + 256;
constexpr int QROPE_OFFSET = QLAT_OFFSET + 8192;
constexpr int MLAT_OFFSET = QROPE_OFFSET + 1024;
constexpr int ATT_OFFSET = MLAT_OFFSET + 16384;

__device__ __forceinline__ float bfround(float value) {
    return __bfloat162float(__float2bfloat16_rn(value));
}

__device__ __forceinline__ __nv_bfloat162 bf2bits(unsigned int bits) {
    __nv_bfloat162_raw value;
    value.x = bits & 0xffffu;
    value.y = bits >> 16;
    return __nv_bfloat162(value);
}

__device__ __forceinline__ __nv_bfloat162 dequant_packed(
    unsigned int bits, __nv_bfloat162 zero, __nv_bfloat162 scale) {
    __nv_bfloat162 value = bf2bits((bits & 0x000f000fu) | 0x43004300u);
    value = __hsub2_rn(value, bf2bits(0x43004300u));
    value = __hsub2_rn(value, zero);
    return __hmul2_rn(value, scale);
}

__device__ __forceinline__ float2 dequant_pair(
    unsigned int bits, __nv_bfloat162 zero, __nv_bfloat162 scale) {
    return __bfloat1622float2(dequant_packed(bits, zero, scale));
}

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

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

__device__ __forceinline__ float block_sum(float value, float* shared) {
    value = warp_sum(value);
    if ((threadIdx.x & 31) == 0) shared[threadIdx.x >> 5] = value;
    __syncthreads();
    if (threadIdx.x < 32) {
        value = threadIdx.x < 8 ? shared[threadIdx.x] : 0.0f;
        value = warp_sum(value);
        if (threadIdx.x == 0) shared[0] = value;
    }
    __syncthreads();
    return shared[0];
}

__device__ __forceinline__ void normalize(
    const float* hidden, const BFloat* input, const BFloat* weight,
    float* output, float* shared) {
    float square = 0.0f;
    float values[9];
    #pragma unroll
    for (int index = 0; index < 9; ++index) {
        int channel = threadIdx.x + index * 256;
        float value = input ? __bfloat162float(input[channel]) : hidden[channel];
        values[index] = value;
        square += value * value;
    }
    float inverse = rsqrtf(block_sum(square, shared) / 2304.0f + 1.0e-6f);
    #pragma unroll
    for (int index = 0; index < 9; ++index) {
        int channel = threadIdx.x + index * 256;
        output[channel] = bfround(values[index] * inverse * __bfloat162float(weight[channel]));
    }
}

template<int Inputs, int Outputs>
__device__ __forceinline__ void gemv(
    Quant quant, const float* activation, float* partial, int task,
    int expert, float* shared, int activation_offset = 0) {
    constexpr int Tiles = (Outputs + 127) / 128;
    int split = task / Tiles;
    int column = (task % Tiles) * 128 + (threadIdx.x & 31) * 4;
    int warp = threadIdx.x >> 5;
    float result[4] = {0, 0, 0, 0};
    const unsigned char* packed = quant.weight + (size_t)expert * (Inputs / 2) * Outputs;
    const BFloat* scales = quant.scales + expert * (Inputs / 128) * Outputs;
    const BFloat* zeros = quant.zeros + expert * (Inputs / 128) * Outputs;
    #pragma unroll
    for (int group = 0; group < 2; ++group) {
        int group_index = split * 2 + group;
        uint2 packed_scale = column < Outputs ? *reinterpret_cast<const uint2*>(scales + group_index * Outputs + column) : make_uint2(0, 0);
        uint2 packed_zero = column < Outputs ? *reinterpret_cast<const uint2*>(zeros + group_index * Outputs + column) : make_uint2(0, 0);
        __nv_bfloat162 scale[2] = {bf2bits(__byte_perm(packed_scale.x, packed_scale.y, 0x5410)), bf2bits(__byte_perm(packed_scale.x, packed_scale.y, 0x7632))};
        __nv_bfloat162 zero[2] = {bf2bits(__byte_perm(packed_zero.x, packed_zero.y, 0x5410)), bf2bits(__byte_perm(packed_zero.x, packed_zero.y, 0x7632))};
        #pragma unroll
        for (int offset = 0; offset < 8; ++offset) {
            int input_pair = group_index * 64 + warp + offset * 8;
            unsigned int bits = column < Outputs ? *reinterpret_cast<const unsigned int*>(packed + input_pair * Outputs + column) : 0;
            float even = activation[input_pair * 2 - activation_offset];
            float odd = activation[input_pair * 2 + 1 - activation_offset];
            #pragma unroll
            for (int component = 0; component < 2; ++component) {
                float2 low = dequant_pair(bits >> (component * 8), zero[component], scale[component]);
                float2 high = dequant_pair(bits >> (component * 8 + 4), zero[component], scale[component]);
                result[component] = fmaf(even, low.x, result[component]);
                result[component + 2] = fmaf(even, low.y, result[component + 2]);
                result[component] = fmaf(odd, high.x, result[component]);
                result[component + 2] = fmaf(odd, high.y, result[component + 2]);
            }
        }
    }
    reinterpret_cast<float4*>(shared)[threadIdx.x] = make_float4(result[0], result[1], result[2], result[3]);
    __syncthreads();
    if (threadIdx.x < 128) {
        float sum = 0;
        #pragma unroll
        for (int slice = 0; slice < 8; ++slice) sum += shared[slice * 128 + threadIdx.x];
        int output = (task % Tiles) * 128 + threadIdx.x;
        if (output < Outputs) partial[split * Outputs + output] = sum;
    }
    __syncthreads();
}

template<int Outputs>
__device__ __forceinline__ void dense_rows(
    const BFloat* weight, const float* input, float* result, int task, float* shared) {
    float accum = 0;
    #pragma unroll
    for (int index = 0; index < 9; ++index) {
        int column = threadIdx.x + index * 256;
        accum = fmaf(__bfloat162float(weight[task * 2304 + column]), input[column], accum);
    }
    accum = block_sum(accum, shared);
    if (threadIdx.x == 0) result[task] = bfround(accum);
}

__device__ __forceinline__ void kda_state(
    const Layer& layer, const State& state, const float* partial,
    const float* beta, float* output, int task, float* shared) {
    int head = task / 8;
    int value_start = (task % 8) * 16;
    int lane = threadIdx.x & 15;
    int warp = threadIdx.x >> 4;
    if (threadIdx.x < 128) {
        int channel = head * 128 + threadIdx.x;
        float projections[4] = {0, 0, 0, 0};
        #pragma unroll
        for (int projection = 0; projection < 4; ++projection) {
            #pragma unroll
            for (int split = 0; split < 9; ++split)
                projections[projection] += partial[(projection * 9 + split) * 4096 + channel];
            projections[projection] = bfround(projections[projection]);
        }
        #pragma unroll
        for (int projection = 0; projection < 3; ++projection) {
            float convolved = 0;
            #pragma unroll
            for (int time = 0; time < 3; ++time) {
                float value = __bfloat162float(state.old_conv[projection][time * 4096 + channel]);
                convolved += value * __bfloat162float(layer.conv[(projection * 4096 + channel) * 4 + time]);
                if (task % 8 == 0 && time > 0)
                    state.new_conv[projection][(time - 1) * 4096 + channel] = __float2bfloat16_rn(value);
            }
            convolved += projections[projection] * __bfloat162float(layer.conv[(projection * 4096 + channel) * 4 + 3]);
            if (task % 8 == 0)
                state.new_conv[projection][2 * 4096 + channel] = __float2bfloat16_rn(projections[projection]);
            shared[projection * 128 + threadIdx.x] = bfround(convolved / (1.0f + __expf(-convolved)));
        }
        shared[384 + threadIdx.x] = 1.0f / (1.0f + __expf(projections[3]));
    }
    __syncthreads();
    float state_values[8];
    float prediction = 0;
    #pragma unroll
    for (int index = 0; index < 8; ++index) {
        int key = warp + index * 16;
        float value = state.old_state[head * 16384 + key * 128 + value_start + lane] * shared[384 + key];
        state_values[index] = value;
        prediction += value * shared[128 + key];
    }
    shared[512 + threadIdx.x] = prediction;
    __syncthreads();
    if (warp == 0) {
        float prediction_sum = 0;
        #pragma unroll
        for (int slice = 0; slice < 16; ++slice) prediction_sum += shared[512 + slice * 16 + lane];
        shared[768 + lane] = (shared[256 + value_start + lane] - prediction_sum) / (1.0f + __expf(-beta[head]));
    }
    __syncthreads();
    float delta = shared[768 + lane];
    float result = 0;
    #pragma unroll
    for (int index = 0; index < 8; ++index) {
        int key = warp + index * 16;
        float value = state_values[index] + shared[128 + key] * delta;
        state.new_state[head * 16384 + key * 128 + value_start + lane] = value;
        result += value * (shared[key] * 0.08838834764831845f);
    }
    shared[512 + threadIdx.x] = result;
    __syncthreads();
    if (warp == 0) {
        float sum = 0;
        #pragma unroll
        for (int slice = 0; slice < 16; ++slice) sum += shared[512 + slice * 16 + lane];
        output[head * 128 + value_start + lane] = bfround(sum);
    }
    __syncthreads();
}

__device__ __forceinline__ void choose_experts(float* route, float scale, int experts) {
    if (threadIdx.x < 32) {
        int lane = threadIdx.x;
        int* indices = reinterpret_cast<int*>(route + 128);
        unsigned int first = __float_as_uint(route[lane]);
        unsigned int second = __float_as_uint(route[lane + 32]);
        first = first & 0x80000000u ? ~first : first ^ 0x80000000u;
        second = second & 0x80000000u ? ~second : second ^ 0x80000000u;
        float total = 0, maximum = 0, chosen = 0;
        int chosen_index = 0;
        #pragma unroll
        for (int rank = 0; rank < 8; ++rank) {
            unsigned int key = __reduce_max_sync(0xffffffff, max(first, second));
            int index = first == key ? lane : second == key ? lane + 32 : 64;
            index = __reduce_min_sync(0xffffffff, index);
            float largest = __uint_as_float(key & 0x80000000u ? key ^ 0x80000000u : ~key);
            if (rank == 0) maximum = largest;
            float probability = __expf(largest - maximum);
            total += probability;
            if (lane == rank) {
                chosen = probability;
                chosen_index = index;
            }
            if (index == lane) first = 0;
            if (index == lane + 32) second = 0;
        }
        if (lane < 8) {
            route[64 + lane] = chosen * (scale / total);
            indices[lane] = chosen_index;
        } else if (lane == 8) {
            indices[8] = 0;
            route[72] = 1.0f;
        }
    }
}

__device__ __forceinline__ void mla_prepare(
    const Arguments& args, const float* partial, float* query,
    BFloat* query_rope) {
    int global = blockIdx.x * 256 + threadIdx.x;
    int stride = gridDim.x * 256;
    for (int pair = global; pair < 3072; pair += stride) {
        int head = pair / 96;
        int dim = (pair % 96) * 2;
        float even = 0, odd = 0;
        #pragma unroll
        for (int split = 0; split < 9; ++split) {
            even += partial[split * 6144 + head * 192 + dim];
            odd += partial[split * 6144 + head * 192 + dim + 1];
        }
        even = bfround(even);
        odd = bfround(odd);
        if (dim < 128) {
            query[head * 192 + dim] = even;
            query[head * 192 + dim + 1] = odd;
        } else {
            int rotary = dim - 128;
            float angle = args.context * powf(10000.0f, -float(rotary) / 64.0f);
            float sine, cosine;
            __sincosf(angle, &sine, &cosine);
            query_rope[head * 64 + rotary] = __float2bfloat16_rn(even * cosine - odd * sine);
            query_rope[head * 64 + rotary + 1] = __float2bfloat16_rn(odd * cosine + even * sine);
        }
    }
    for (int pair = global; pair < 288; pair += stride) {
        float even = 0, odd = 0;
        #pragma unroll
        for (int split = 0; split < 9; ++split) {
            even += partial[9 * 6144 + split * 576 + pair * 2];
            odd += partial[9 * 6144 + split * 576 + pair * 2 + 1];
        }
        even = bfround(even);
        odd = bfround(odd);
        if (pair < 256) {
            args.latent[args.context * 512 + pair * 2] = __float2bfloat16_rn(even);
            args.latent[args.context * 512 + pair * 2 + 1] = __float2bfloat16_rn(odd);
        } else {
            int rotary = (pair - 256) * 2;
            float angle = args.context * powf(10000.0f, -float(rotary) / 64.0f);
            float sine, cosine;
            __sincosf(angle, &sine, &cosine);
            args.rope[args.context * 64 + rotary] = __float2bfloat16_rn(even * cosine - odd * sine);
            args.rope[args.context * 64 + rotary + 1] = __float2bfloat16_rn(odd * cosine + even * sine);
        }
    }
    if (args.old_latent != args.latent)
        for (int index = global; index < args.context * 512; index += stride)
            args.latent[index] = args.old_latent[index];
    if (args.old_rope != args.rope)
        for (int index = global; index < args.context * 64; index += stride)
            args.rope[index] = args.old_rope[index];
    int padded = ((args.context + 64) / 64) * 64;
    for (int index = (args.context + 1) * 512 + global; index < padded * 512; index += stride)
        args.latent[index] = __float2bfloat16_rn(0);
    for (int index = (args.context + 1) * 64 + global; index < padded * 64; index += stride)
        args.rope[index] = __float2bfloat16_rn(0);
}

__device__ __forceinline__ void absorb_query(
    Quant quant, const float* query, BFloat* absorbed, int task) {
    int lane = threadIdx.x & 31;
    int latent = (task % 64) * 8 + (threadIdx.x >> 5);
    int head = task / 64;
    float result = 0;
    #pragma unroll
    for (int part = 0; part < 4; ++part) {
        int channel = head * 256 + lane + part * 32;
        unsigned int bits = quant.weight[(latent / 2) * 8192 + channel];
        float value = float((bits >> ((latent & 1) * 4)) & 15);
        float zero = __bfloat162float(quant.zeros[(latent / 128) * 8192 + channel]);
        float scale = __bfloat162float(quant.scales[(latent / 128) * 8192 + channel]);
        value = bfround(bfround(value - zero) * scale);
        result = fmaf(value, query[head * 192 + lane + part * 32], result);
    }
    result = warp_sum(result);
    if (lane == 0) absorbed[head * 512 + latent] = __float2bfloat16_rn(result);
}

__device__ __forceinline__ void async_latent_tile(
    const BFloat* latent, BFloat* destination, int task, int feature) {
    #pragma unroll
    for (int vector = threadIdx.x; vector < 1024; vector += 256) {
        int row = vector / 16;
        int column = (vector % 16) * 8;
        unsigned int address = __cvta_generic_to_shared(destination + row * 136 + column);
        const BFloat* source = latent + (task * 64 + row) * 512 + feature + column;
        asm volatile("cp.async.ca.shared.global [%0], [%1], 16;" :: "r"(address), "l"(source) : "memory");
    }
    asm volatile("cp.async.commit_group;" ::: "memory");
}

__device__ __forceinline__ void latent_attention(
    const Arguments& args, const BFloat* query, const BFloat* query_rope,
    float* results, float* stats, int task, float* shared) {
    int warp = threadIdx.x >> 5;
    int lane = threadIdx.x & 31;
    int head = (warp / 4) * 16;
    wm::fragment<wm::matrix_a, 16, 16, 16, BFloat, wm::row_major> query_fragment;
    wm::fragment<wm::matrix_b, 16, 16, 16, BFloat, wm::col_major> key_fragment;
    wm::fragment<wm::accumulator, 16, 16, 16, float> accumulator;
    BFloat* shared_query = reinterpret_cast<BFloat*>(shared);
    BFloat* shared_key = reinterpret_cast<BFloat*>(shared + 3584);
    wm::fill_fragment(accumulator, 0.0f);
    async_latent_tile(args.latent, shared_key, task, 0);
    #pragma unroll 1
    for (int feature = 0; feature < 512; feature += 128) {
        for (int vector = threadIdx.x; vector < 512; vector += 256) {
            int row = vector / 16;
            int column = (vector % 16) * 8;
            *reinterpret_cast<uint4*>(shared_query + row * 136 + column) = *reinterpret_cast<const uint4*>(query + row * 512 + feature + column);
        }
        BFloat* current_key = shared_key + ((feature / 128) % 2) * 8704;
        if (feature < 384) {
            async_latent_tile(args.latent, shared_key + (((feature / 128) + 1) % 2) * 8704, task, feature + 128);
            asm volatile("cp.async.wait_group 1;" ::: "memory");
        } else {
            asm volatile("cp.async.wait_group 0;" ::: "memory");
        }
        __syncthreads();
        #pragma unroll
        for (int dim = 0; dim < 128; dim += 16) {
            wm::load_matrix_sync(query_fragment, shared_query + head * 136 + dim, 136);
            wm::load_matrix_sync(key_fragment, current_key + (warp % 4) * 16 * 136 + dim, 136);
            wm::mma_sync(accumulator, query_fragment, key_fragment, accumulator);
        }
        __syncthreads();
    }
    for (int vector = threadIdx.x; vector < 256; vector += 256)
        *reinterpret_cast<uint4*>(shared_query + (vector / 8) * 72 + (vector % 8) * 8) = reinterpret_cast<const uint4*>(query_rope)[vector];
    for (int vector = threadIdx.x; vector < 512; vector += 256)
        *reinterpret_cast<uint4*>(shared_key + (vector / 8) * 72 + (vector % 8) * 8) = reinterpret_cast<const uint4*>(args.rope + task * 64 * 64)[vector];
    __syncthreads();
    #pragma unroll
    for (int dim = 0; dim < 64; dim += 16) {
        wm::load_matrix_sync(query_fragment, shared_query + head * 72 + dim, 72);
        wm::load_matrix_sync(key_fragment, shared_key + (warp % 4) * 16 * 72 + dim, 72);
        wm::mma_sync(accumulator, query_fragment, key_fragment, accumulator);
    }
    __syncthreads();
    wm::store_matrix_sync(shared + head * 72 + (warp % 4) * 16, accumulator, 72, wm::mem_row_major);
    __syncthreads();
    BFloat* probabilities = reinterpret_cast<BFloat*>(shared + 2304);
    #pragma unroll
    for (int iteration = 0; iteration < 4; ++iteration) {
        int row = warp + iteration * 8;
        float first = task * 64 + lane <= args.context ? shared[row * 72 + lane] * 0.07216878364870322f : -INFINITY;
        float second = task * 64 + lane + 32 <= args.context ? shared[row * 72 + lane + 32] * 0.07216878364870322f : -INFINITY;
        float maximum = warp_max(fmaxf(first, second));
        maximum = __shfl_sync(0xffffffff, maximum, 0);
        first = __expf(first - maximum);
        second = __expf(second - maximum);
        float sum = warp_sum(first + second);
        probabilities[row * 72 + lane] = __float2bfloat16_rn(first);
        probabilities[row * 72 + lane + 32] = __float2bfloat16_rn(second);
        if (lane == 0) {
            stats[task * 64 + row] = maximum;
            stats[task * 64 + 32 + row] = sum;
        }
    }
    __syncthreads();
    wm::fragment<wm::matrix_b, 16, 16, 16, BFloat, wm::row_major> value_fragment;
    async_latent_tile(args.latent, shared_key, task, 0);
    #pragma unroll 1
    for (int feature = 0; feature < 512; feature += 128) {
        BFloat* current_key = shared_key + ((feature / 128) % 2) * 8704;
        if (feature < 384) {
            async_latent_tile(args.latent, shared_key + (((feature / 128) + 1) % 2) * 8704, task, feature + 128);
            asm volatile("cp.async.wait_group 1;" ::: "memory");
        } else {
            asm volatile("cp.async.wait_group 0;" ::: "memory");
        }
        __syncthreads();
        #pragma unroll
        for (int tile = warp; tile < 16; tile += 8) {
            int row = (tile / 8) * 16;
            int dim = (tile % 8) * 16;
            wm::fill_fragment(accumulator, 0.0f);
            #pragma unroll
            for (int inner = 0; inner < 64; inner += 16) {
                wm::load_matrix_sync(query_fragment, probabilities + row * 72 + inner, 72);
                wm::load_matrix_sync(value_fragment, current_key + inner * 136 + dim, 136);
                wm::mma_sync(accumulator, query_fragment, value_fragment, accumulator);
            }
            wm::store_matrix_sync(results + task * 16384 + row * 512 + feature + dim, accumulator, 512, wm::mem_row_major);
        }
        __syncthreads();
    }
    __syncthreads();
}

__device__ __forceinline__ void merge_attention(
    const float* partial, const float* stats, float* latent, int parts,
    int task, float* shared) {
    int head = task / 8;
    float maximum = -INFINITY;
    for (int part = threadIdx.x; part < parts; part += 256)
        maximum = fmaxf(maximum, stats[part * 64 + head]);
    maximum = warp_max(maximum);
    if ((threadIdx.x & 31) == 0) shared[threadIdx.x >> 5] = maximum;
    __syncthreads();
    if (threadIdx.x < 32) {
        maximum = threadIdx.x < 8 ? shared[threadIdx.x] : -INFINITY;
        maximum = warp_max(maximum);
        if (threadIdx.x == 0) shared[0] = maximum;
    }
    __syncthreads();
    maximum = shared[0];
    float denominator = 0;
    for (int part = threadIdx.x; part < parts; part += 256) {
        float factor = __expf(stats[part * 64 + head] - maximum);
        shared[256 + part] = factor;
        denominator += factor * stats[part * 64 + 32 + head];
    }
    denominator = block_sum(denominator, shared);
    int dim = (task % 8) * 64 + (threadIdx.x % 64);
    float result = 0;
    for (int part = threadIdx.x / 64; part < parts; part += 4)
        result = fmaf(shared[256 + part], partial[part * 16384 + head * 512 + dim], result);
    shared[7936 + threadIdx.x] = result;
    __syncthreads();
    if (threadIdx.x < 64) {
        float sum = shared[7936 + threadIdx.x] + shared[8000 + threadIdx.x] + shared[8064 + threadIdx.x] + shared[8128 + threadIdx.x];
        latent[head * 512 + dim] = sum / denominator;
    }
    __syncthreads();
}

__device__ __forceinline__ void latent_values(
    Quant quant, const float* latent, float* output, int task, float* shared) {
    int head = task / 4;
    int group = task % 4;
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    float result[4] = {0, 0, 0, 0};
    int channel = head * 256 + 128 + lane * 4;
    uint2 packed_scale = *reinterpret_cast<const uint2*>(quant.scales + group * 8192 + channel);
    uint2 packed_zero = *reinterpret_cast<const uint2*>(quant.zeros + group * 8192 + channel);
    __nv_bfloat162 scale[2] = {bf2bits(__byte_perm(packed_scale.x, packed_scale.y, 0x5410)), bf2bits(__byte_perm(packed_scale.x, packed_scale.y, 0x7632))};
    __nv_bfloat162 zero[2] = {bf2bits(__byte_perm(packed_zero.x, packed_zero.y, 0x5410)), bf2bits(__byte_perm(packed_zero.x, packed_zero.y, 0x7632))};
    #pragma unroll
    for (int offset = 0; offset < 8; ++offset) {
        int pair = group * 64 + warp + offset * 8;
        unsigned int bits = *reinterpret_cast<const unsigned int*>(quant.weight + pair * 8192 + channel);
        float even = latent[head * 512 + pair * 2];
        float odd = latent[head * 512 + pair * 2 + 1];
        #pragma unroll
        for (int component = 0; component < 2; ++component) {
            float2 low = dequant_pair(bits >> (component * 8), zero[component], scale[component]);
            float2 high = dequant_pair(bits >> (component * 8 + 4), zero[component], scale[component]);
            result[component] = fmaf(even, low.x, result[component]);
            result[component + 2] = fmaf(even, low.y, result[component + 2]);
            result[component] = fmaf(odd, high.x, result[component]);
            result[component + 2] = fmaf(odd, high.y, result[component + 2]);
        }
    }
    reinterpret_cast<float4*>(shared)[threadIdx.x] = make_float4(result[0], result[1], result[2], result[3]);
    __syncthreads();
    if (threadIdx.x < 128) {
        float sum = 0;
        #pragma unroll
        for (int slice = 0; slice < 8; ++slice) sum += shared[slice * 128 + threadIdx.x];
        output[group * 4096 + head * 128 + threadIdx.x] = sum;
    }
    __syncthreads();
}

__global__ __launch_bounds__(256, 2) void hybrid_decode(Arguments args) {
    __shared__ __align__(32) float shared[12288];
    cg::grid_group grid = cg::this_grid();
    float* normalized = args.workspace + X_OFFSET;
    float* hidden = args.workspace + H_OFFSET;
    float* partial = args.workspace + PART_OFFSET;
    float* aux = args.workspace + AUX_OFFSET;
    float* route = args.workspace + ROUTE_OFFSET;
    BFloat* query_latent = reinterpret_cast<BFloat*>(args.workspace + QLAT_OFFSET);
    BFloat* query_rope = reinterpret_cast<BFloat*>(args.workspace + QROPE_OFFSET);
    float* mixed_latent = args.workspace + MLAT_OFFSET;
    float* att_partial = args.workspace + ATT_OFFSET;
    int attention_parts = (args.context + 64) / 64;
    float* att_stats = att_partial + attention_parts * 16384;
    int global = blockIdx.x * 256 + threadIdx.x;
    int stride = gridDim.x * 256;

    for (int layer_index = 0; layer_index < 4; ++layer_index) {
        const Layer& layer = args.layers[layer_index];
        float* local_normalized = shared + 1024;
        normalize(hidden, layer_index == 0 ? args.hidden : nullptr, layer.attn_norm, local_normalized, shared);
        __syncthreads();
        if (layer_index < 3) {
            for (int task = blockIdx.x; task < 4 * 9 * 32; task += gridDim.x) {
                int projection = task / (9 * 32);
                gemv<2304, 4096>(layer.attn[projection], local_normalized,
                    partial + projection * 9 * 4096, task % (9 * 32), 0, shared);
            }
            if (blockIdx.x < 32) dense_rows<32>(layer.beta, local_normalized, route, blockIdx.x, shared);
            grid.sync();
            for (int task = blockIdx.x; task < 256; task += gridDim.x)
                kda_state(layer, args.states[layer_index], partial, route, aux, task, shared);
            grid.sync();
        } else {
            for (int task = blockIdx.x; task < 9 * 48 + 9 * 5; task += gridDim.x) {
                if (task < 9 * 48)
                    gemv<2304, 6144>(layer.attn[0], local_normalized, partial, task, 0, shared);
                else
                    gemv<2304, 576>(layer.attn[1], local_normalized, partial + 9 * 6144, task - 9 * 48, 0, shared);
            }
            grid.sync();
            mla_prepare(args, partial, aux, query_rope);
            grid.sync();
            for (int task = blockIdx.x; task < 32 * 64; task += gridDim.x)
                absorb_query(layer.attn[2], aux, query_latent, task);
            grid.sync();
            for (int task = blockIdx.x; task < attention_parts; task += gridDim.x)
                latent_attention(args, query_latent, query_rope, att_partial, att_stats, task, shared);
            grid.sync();
            for (int task = blockIdx.x; task < 256; task += gridDim.x)
                merge_attention(att_partial, att_stats, mixed_latent, attention_parts, task, shared);
            grid.sync();
            for (int task = blockIdx.x; task < 128; task += gridDim.x)
                latent_values(layer.attn[2], mixed_latent, aux, task, shared);
            grid.sync();
        }
        for (int task = blockIdx.x; task < 16 * 18; task += gridDim.x) {
            if (layer_index == 3) {
                int offset = (task / 18) * 256;
                int channel = offset + threadIdx.x;
                local_normalized[threadIdx.x] = bfround(aux[channel] + aux[4096 + channel] + aux[8192 + channel] + aux[12288 + channel]);
                __syncthreads();
                gemv<4096, 2304>(layer.attn[4], local_normalized, partial, task, 0, shared, offset);
            } else {
                gemv<4096, 2304>(layer.attn[4], aux, partial, task, 0, shared);
            }
        }
        grid.sync();
        float* post_attention = aux + 4096;
        if (blockIdx.x < args.experts) {
            #pragma unroll
            for (int index = 0; index < 9; ++index) {
                int channel = threadIdx.x + index * 256;
                float value = 0;
                #pragma unroll
                for (int split = 0; split < 16; ++split) value += partial[split * 2304 + channel];
                float residual = layer_index == 0 ? __bfloat162float(args.hidden[channel]) : hidden[channel];
                float result = bfround(residual + bfround(value));
                local_normalized[channel] = result;
                if (blockIdx.x == 0) post_attention[channel] = result;
            }
            __syncthreads();
            normalize(local_normalized, nullptr, layer.moe_norm, local_normalized, shared);
            __syncthreads();
            dense_rows<64>(layer.router, local_normalized, route, blockIdx.x, shared);
            if (blockIdx.x == 0)
                for (int channel = threadIdx.x; channel < 2304; channel += 256)
                    normalized[channel] = local_normalized[channel];
        }
        grid.sync();
        float* local_route = shared + 3328;
        if (threadIdx.x < 64) local_route[threadIdx.x] = route[threadIdx.x];
        __syncthreads();
        choose_experts(local_route, args.routing_scale, args.experts);
        __syncthreads();
        if (blockIdx.x == 0 && threadIdx.x < 9)
            route[64 + threadIdx.x] = local_route[64 + threadIdx.x];
        const int* indices = reinterpret_cast<const int*>(local_route + 128);
        for (int task = blockIdx.x; task < 2 * 9 * 9 * 8; task += gridDim.x) {
            int projection = task / (9 * 9 * 8);
            int expert = (task / (9 * 8)) % 9;
            int quant_index = projection + (expert == 8 ? 3 : 0);
            gemv<2304, 1024>(layer.moe[quant_index], normalized,
                partial + (projection * 9 + expert) * 9 * 1024,
                task % (9 * 8), indices[expert], shared);
        }
        grid.sync();
        for (int task = blockIdx.x; task < 9 * 4 * 18; task += gridDim.x) {
            int expert = task / (4 * 18);
            int offset = ((task / 18) % 4) * 256;
            int channel = offset + threadIdx.x;
            float gate = 0, up = 0;
            #pragma unroll
            for (int split = 0; split < 9; ++split) {
                gate += partial[(expert * 9 + split) * 1024 + channel];
                up += partial[((9 + expert) * 9 + split) * 1024 + channel];
            }
            local_normalized[threadIdx.x] = (gate / (1.0f + __expf(-gate))) * up;
            __syncthreads();
            gemv<1024, 2304>(layer.moe[expert == 8 ? 5 : 2], local_normalized,
                att_partial + expert * 4 * 2304, task % (4 * 18), indices[expert], shared, offset);
        }
        grid.sync();
        for (int channel = global; channel < 2304; channel += stride) {
            float value = 0;
            #pragma unroll
            for (int expert = 0; expert < 9; ++expert) {
                float output = 0;
                #pragma unroll
                for (int split = 0; split < 4; ++split)
                    output += att_partial[(expert * 4 + split) * 2304 + channel];
                value += route[64 + expert] * output;
            }
            float result = bfround(post_attention[channel] + bfround(value));
            if (layer_index == 3) args.output[channel] = __float2bfloat16_rn(result);
            else hidden[channel] = result;
        }
        if (layer_index < 3) grid.sync();
    }
}

class Runner {
    Arguments args_{};
    std::vector<torch::Tensor> weights_;
    int blocks_;
public:
    Runner(std::vector<torch::Tensor> weights, int experts, double routing_scale)
        : weights_(std::move(weights)) {
        size_t position = 0;
        auto next_bf = [&]() { return reinterpret_cast<const BFloat*>(weights_[position++].data_ptr()); };
        auto next_quant = [&]() {
            Quant quant;
            quant.weight = weights_[position++].data_ptr<unsigned char>();
            quant.scales = next_bf();
            quant.zeros = next_bf();
            return quant;
        };
        for (int index = 0; index < 4; ++index) {
            Layer& layer = args_.layers[index];
            layer.attn_norm = next_bf();
            layer.moe_norm = next_bf();
            if (index < 3) {
                for (int projection = 0; projection < 4; ++projection) layer.attn[projection] = next_quant();
                layer.beta = next_bf();
                layer.conv = next_bf();
            } else {
                for (int projection = 0; projection < 3; ++projection) layer.attn[projection] = next_quant();
            }
            layer.attn[4] = next_quant();
            layer.router = next_bf();
            for (int projection = 0; projection < 6; ++projection) layer.moe[projection] = next_quant();
        }
        args_.experts = experts;
        args_.routing_scale = float(routing_scale);
        int device, multiprocessors, resident;
        C10_CUDA_CHECK(cudaGetDevice(&device));
        C10_CUDA_CHECK(cudaDeviceGetAttribute(&multiprocessors, cudaDevAttrMultiProcessorCount, device));
        C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&resident, hybrid_decode, 256, 0));
        blocks_ = multiprocessors * std::min(resident, 2);
    }

    void step(torch::Tensor hidden, torch::Tensor output, torch::Tensor workspace,
              std::vector<torch::Tensor> old_state, std::vector<torch::Tensor> new_state,
              int context) {
        Arguments args = args_;
        args.hidden = reinterpret_cast<const BFloat*>(hidden.data_ptr());
        args.output = reinterpret_cast<BFloat*>(output.data_ptr());
        args.workspace = workspace.data_ptr<float>();
        args.context = context;
        for (int layer = 0; layer < 3; ++layer) {
            args.states[layer].old_state = old_state[layer * 4].data_ptr<float>();
            args.states[layer].new_state = new_state[layer * 4].data_ptr<float>();
            for (int projection = 0; projection < 3; ++projection) {
                args.states[layer].old_conv[projection] = reinterpret_cast<const BFloat*>(old_state[layer * 4 + 1 + projection].data_ptr());
                args.states[layer].new_conv[projection] = reinterpret_cast<BFloat*>(new_state[layer * 4 + 1 + projection].data_ptr());
            }
        }
        args.old_latent = reinterpret_cast<const BFloat*>(old_state[12].data_ptr());
        args.old_rope = reinterpret_cast<const BFloat*>(old_state[13].data_ptr());
        args.latent = reinterpret_cast<BFloat*>(new_state[12].data_ptr());
        args.rope = reinterpret_cast<BFloat*>(new_state[13].data_ptr());
        void* parameters[] = {&args};
        auto stream = at::cuda::getCurrentCUDAStream();
        C10_CUDA_CHECK(cudaLaunchCooperativeKernel(reinterpret_cast<void*>(hybrid_decode),
            dim3(blocks_), dim3(256), parameters, 0, stream));
    }

    int blocks() const { return blocks_; }
};

PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
    pybind11::class_<Runner>(module, "Runner")
        .def(pybind11::init<std::vector<torch::Tensor>, int, double>())
        .def("step", &Runner::step)
        .def("blocks", &Runner::blocks);
}
"""

_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        os.environ["PATH"] = ninja.BIN_DIR + os.pathsep + os.environ["PATH"]
        _extension = load_inline(
            name="kimi_hybrid_megakernel",
            cpp_sources="",
            cuda_sources=_CUDA,
            extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo", "-gencode=arch=compute_120,code=sm_120"],
            extra_cflags=["-O3"],
            verbose=False,
        )
    return _extension


class Model(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        if (cfg.hidden, cfg.kda_heads, cfg.kda_head_dim, cfg.short_conv,
            cfg.mla_heads, cfg.kv_lora, cfg.qk_nope, cfg.qk_rope, cfg.v_head,
            cfg.n_experts, cfg.n_active, cfg.n_shared, cfg.moe_inter, cfg.group,
            tuple(cfg.pattern), cfg.dtype, cfg.rope_theta) != (
                2304, 32, 128, 4, 32, 512, 128, 64, 128,
                64, 8, 1, 1024, 128, ("K", "K", "K", "M"), torch.bfloat16, 10000.0):
            raise ValueError("This kernel specializes the Kimi-Linear benchmark motif")
        self.cfg = cfg
        self.blocks = nn.ModuleList(Block(cfg, kind) for kind in cfg.pattern)
        self._runner = None
        self._scratch = None

    def _prepare(self):
        weights = []

        def quant(module):
            weights.extend((module.w_q, module.scales, module.zeros))

        for index, block in enumerate(self.blocks):
            weights.extend((block.attn_norm, block.moe_norm))
            if index < 3:
                for name in ("q_proj", "k_proj", "v_proj", "g_proj"):
                    quant(getattr(block.attn, name))
                weights.extend((block.attn.beta_proj.weight, block.attn.conv_w))
            else:
                for name in ("q_proj", "kv_a", "kv_b"):
                    quant(getattr(block.attn, name))
            quant(block.attn.o_proj)
            weights.append(block.moe.router.weight)
            for name in ("gate", "up", "down", "s_gate", "s_up", "s_down"):
                quant(getattr(block.moe, name))
        self._runner = _get_extension().Runner(weights, self.cfg.n_experts, self.cfg.routed_scaling)

    def step(self, hidden, state):
        if self._runner is None:
            self._prepare()
        context = state[3]["c_kv"].shape[0]
        parts = (context + 64) // 64
        required = 251136 + max(parts * (16384 + 64), 9 * 4 * 2304)
        if self._scratch is None or self._scratch.numel() < required:
            self._scratch = torch.empty(required + 1048576, device=hidden.device, dtype=torch.float32)
        output = torch.empty_like(hidden)
        old_tensors = []
        new_tensors = []
        for index in range(3):
            block_state = state[index]
            for key in ("S", "cq", "ck", "cv"):
                previous = block_state[key]
                updated = torch.empty_like(previous)
                old_tensors.append(previous)
                new_tensors.append(updated)
                block_state[key] = updated
        for key, width in (("c_kv", 512), ("k_rope", 64)):
            previous = state[3][key]
            capacity = previous.untyped_storage().nbytes() // (width * 2)
            if capacity >= parts * 64 and previous.storage_offset() == 0:
                updated = previous.as_strided((context + 1, width), (width, 1))
            else:
                capacity = ((context + 1024) // 1024) * 1024
                updated = torch.empty((capacity, width), device=hidden.device, dtype=hidden.dtype)[:context + 1]
            old_tensors.append(previous)
            new_tensors.append(updated)
            state[3][key] = updated
        self._runner.step(hidden, output, self._scratch, old_tensors, new_tensors, context)
        return output, state


def init_state(cfg, context_len, seed):
    device = torch.device("cuda:0")
    generator = torch.Generator(device=device).manual_seed(seed)
    state = []
    for kind in cfg.pattern:
        if kind == "K":
            block_state = {"S": torch.randn(32, 128, 128, device=device, generator=generator) * 0.05}
            for key in ("cq", "ck", "cv"):
                block_state[key] = torch.randn(3, 4096, device=device, generator=generator, dtype=cfg.dtype) * 0.1
        else:
            block_state = {
                "c_kv": torch.randn(context_len, 512, device=device, generator=generator, dtype=cfg.dtype) * 0.1,
                "k_rope": torch.randn(context_len, 64, device=device, generator=generator, dtype=cfg.dtype) * 0.1,
            }
        state.append(block_state)
    return state


def init_token(cfg, seed):
    device = torch.device("cuda:0")
    generator = torch.Generator(device=device).manual_seed(seed + 1)
    return torch.randn(cfg.hidden, device=device, generator=generator, dtype=cfg.dtype) * 0.25