"""Kimi-Linear W4A16 Hybrid Decode Megakernel for SM120 Blackwell.""" from __future__ import annotations import os from dataclasses import dataclass, field import ninja import torch import torch.nn as nn os.environ["PATH"] = ninja.BIN_DIR + os.pathsep + os.environ.get("PATH", "") from torch.utils.cpp_extension import load_inline OP_TYPE = "kimi_linear_w4a16_decode" HARDWARE_REQUIRED = ["RTX_PRO_6000"] EPS = 1.0e-6 GROUP_SIZE = 128 @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: dict) -> Config: return Config(n_experts=int(shape.get("n_experts", 64))) def _pack_int4(w_q: torch.Tensor) -> torch.Tensor: lo = w_q[0::2] & 0xF hi = w_q[1::2] & 0xF return (lo | (hi << 4)).contiguous() def quantize(w_io: torch.Tensor, group: int = GROUP_SIZE): K, N = w_io.shape ng = K // group wg = w_io.view(ng, group, N).float() wmin = wg.min(dim=1, keepdim=True).values wmax = wg.max(dim=1, keepdim=True).values scales = (wmax - wmin).clamp_min(1e-8) / 15.0 zeros = (-wmin / scales).round().clamp(0, 15) w_q = ((wg / scales) + zeros).round().clamp(0, 15).to(torch.uint8).view(K, N) return _pack_int4(w_q), scales.squeeze(1).to(torch.bfloat16), zeros.squeeze(1).to(torch.bfloat16) class QuantLinear(nn.Module): def __init__(self, in_f: int, out_f: int, group: int = GROUP_SIZE): super().__init__() assert in_f % group == 0 and in_f % 2 == 0 self.in_f, self.out_f, self.group = in_f, out_f, group ng = in_f // group self.register_buffer("w_q", torch.zeros(in_f // 2, out_f, dtype=torch.uint8)) self.register_buffer("scales", torch.zeros(ng, out_f, dtype=torch.bfloat16)) self.register_buffer("zeros", torch.zeros(ng, out_f, dtype=torch.bfloat16)) def init_random(self, gen: torch.Generator, std: float = 0.02) -> None: w = torch.randn(self.in_f, self.out_f, generator=gen) * std wq, s, z = quantize(w, self.group) self.w_q.copy_(wq) self.scales.copy_(s) self.zeros.copy_(z) class QuantExperts(nn.Module): def __init__(self, n: int, in_f: int, out_f: int, group: int = GROUP_SIZE): super().__init__() self.n, self.in_f, self.out_f, self.group = n, in_f, out_f, group ng = in_f // group self.register_buffer("w_q", torch.zeros(n, in_f // 2, out_f, dtype=torch.uint8)) self.register_buffer("scales", torch.zeros(n, ng, out_f, dtype=torch.bfloat16)) self.register_buffer("zeros", torch.zeros(n, ng, out_f, dtype=torch.bfloat16)) def init_random(self, gen: torch.Generator, std: float = 0.02) -> None: for e in range(self.n): w = torch.randn(self.in_f, self.out_f, generator=gen) * std wq, s, z = quantize(w, self.group) self.w_q[e].copy_(wq) self.scales[e].copy_(s) self.zeros[e].copy_(z) class KDA(nn.Module): def __init__(self, cfg: Config): super().__init__() self.cfg = cfg H, Dk, d = cfg.kda_heads, cfg.kda_head_dim, cfg.hidden self.q_proj = QuantLinear(d, H * Dk, cfg.group) self.k_proj = QuantLinear(d, H * Dk, cfg.group) self.v_proj = QuantLinear(d, H * Dk, cfg.group) self.g_proj = QuantLinear(d, H * Dk, cfg.group) self.beta_proj = nn.Linear(d, H, bias=False, dtype=cfg.dtype) self.conv_w = nn.Parameter(torch.zeros(3, H * Dk, cfg.short_conv, dtype=cfg.dtype)) self.o_proj = QuantLinear(H * Dk, d, cfg.group) self.scale = Dk ** -0.5 class MLA(nn.Module): def __init__(self, cfg: Config): super().__init__() self.cfg = cfg H, d = cfg.mla_heads, cfg.hidden self.q_proj = QuantLinear(d, H * (cfg.qk_nope + cfg.qk_rope), cfg.group) self.kv_a = QuantLinear(d, cfg.kv_lora + cfg.qk_rope, cfg.group) self.kv_b = QuantLinear(cfg.kv_lora, H * (cfg.qk_nope + cfg.v_head), cfg.group) self.o_proj = QuantLinear(H * cfg.v_head, d, cfg.group) self.scale = (cfg.qk_nope + cfg.qk_rope) ** -0.5 class MoE(nn.Module): def __init__(self, cfg: Config): super().__init__() self.cfg = cfg d, m, E = cfg.hidden, cfg.moe_inter, cfg.n_experts self.router = nn.Linear(d, E, bias=False, dtype=cfg.dtype) self.gate = QuantExperts(E, d, m, cfg.group) self.up = QuantExperts(E, d, m, cfg.group) self.down = QuantExperts(E, m, d, cfg.group) self.s_gate = QuantExperts(cfg.n_shared, d, m, cfg.group) self.s_up = QuantExperts(cfg.n_shared, d, m, cfg.group) self.s_down = QuantExperts(cfg.n_shared, m, d, cfg.group) class Block(nn.Module): def __init__(self, cfg: Config, kind: str): super().__init__() self.kind = kind 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) # --------------------------------------------------------------------------- # # Single Fused CUDA Megakernel across 188 SMs (376 Persistent Co-Resident Blocks) # --------------------------------------------------------------------------- # _CUDA_SRC = r""" #include #include #include #include #include constexpr int HIDDEN = 2304; constexpr int KDA_H = 32; constexpr int KDA_DK = 128; constexpr int KDA_C = 4096; constexpr int MLA_H = 32; constexpr int KV_LORA = 512; constexpr int QK_NOPE = 128; constexpr int QK_ROPE = 64; constexpr int V_HEAD = 128; constexpr int Q_DIM = 192; constexpr int KVB_DIM = 256; constexpr int MOE_E = 64; constexpr int MOE_ACTIVE = 8; constexpr int MOE_TOTAL_ACTIVE = 9; constexpr int MOE_INTER = 1024; constexpr int GROUP = 128; constexpr int NUM_SMS = 376; constexpr int BLOCK_THREADS = 256; constexpr int NUM_L_SPLITS = 188; __device__ uint32_t g_sync_arr[2] = {0u, 0u}; __device__ __forceinline__ void grid_barrier(uint32_t& local_epoch) { __syncthreads(); if (threadIdx.x == 0) { uint32_t next_epoch = local_epoch + 1u; asm volatile("membar.gl;" ::: "memory"); uint32_t arrived = atomicAdd(&g_sync_arr[0], 1u); if (arrived == NUM_SMS - 1) { *reinterpret_cast(&g_sync_arr[0]) = 0u; asm volatile("membar.gl;" ::: "memory"); *reinterpret_cast(&g_sync_arr[1]) = next_epoch; } else { while (*reinterpret_cast(&g_sync_arr[1]) != next_epoch) { __nanosleep(40); } asm volatile("membar.gl;" ::: "memory"); } local_epoch = next_epoch; } __syncthreads(); } struct QLinearPtrs { const uint8_t* w_q; const __nv_bfloat16* scales; const __nv_bfloat16* zeros; }; struct MoEPtrs { const __nv_bfloat16* router_w; QLinearPtrs gate; QLinearPtrs up; QLinearPtrs down; QLinearPtrs s_gate; QLinearPtrs s_up; QLinearPtrs s_down; }; struct KDABlockPtrs { const __nv_bfloat16* attn_norm; const __nv_bfloat16* moe_norm; QLinearPtrs q_proj; QLinearPtrs k_proj; QLinearPtrs v_proj; QLinearPtrs g_proj; const __nv_bfloat16* beta_w; const __nv_bfloat16* conv_w; QLinearPtrs o_proj; MoEPtrs moe; float* S; __nv_bfloat16* cq; __nv_bfloat16* ck; __nv_bfloat16* cv; }; struct MLABlockPtrs { const __nv_bfloat16* attn_norm; const __nv_bfloat16* moe_norm; QLinearPtrs q_proj; QLinearPtrs kv_a; QLinearPtrs kv_b; QLinearPtrs o_proj; MoEPtrs moe; const __nv_bfloat16* c_kv_src; const __nv_bfloat16* k_rope_src; __nv_bfloat16* c_kv_dst; __nv_bfloat16* k_rope_dst; int pos; }; struct ScratchPtrs { float* norm_x; // [2304] float* attn_out_fp32; // [2304] float* moe_out_fp32; // [2304] __nv_bfloat16* res_h; // [2304] float* kda_qkvg_fp32; // [4 * 4096] float* kda_o; // [4096] float* kda_beta; // [32] float* mla_q_kv_fp32; // [6144 + 576] float* mla_q_abs; // [32 * 512] float* mla_q_rope; // [32 * 64] float* mla_att_m; // [188 * 32] float* mla_att_d; // [188 * 32] float* mla_att_o; // [188 * 32 * 512] float* router_logits; // [64] float* moe_gate_fp32; // [9 * 1024] float* moe_up_fp32; // [9 * 1024] }; struct KernelParams { KDABlockPtrs kda[3]; MLABlockPtrs mla; ScratchPtrs scratch; const __nv_bfloat16* hidden_in; __nv_bfloat16* hidden_out; }; __device__ __forceinline__ float bf16_to_f32(__nv_bfloat16 v) { return __bfloat162float(v); } __device__ __forceinline__ __nv_bfloat16 f32_to_bf16(float v) { return __float2bfloat16(v); } __device__ __forceinline__ float round_bf16(float v) { return __bfloat162float(__float2bfloat16(v)); } __device__ __forceinline__ float warp_reduce_sum(float val) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { val += __shfl_xor_sync(0xffffffff, val, offset); } return val; } __device__ __forceinline__ void accum_u32_fast(uint32_t p, float2 xv, uint32_t magic, float (&g_acc)[4]) { uint32_t p_ev = p & 0x0F0F0F0Fu; uint32_t p_od = (p >> 4) & 0x0F0F0F0Fu; g_acc[0] = fmaf(xv.x, __uint_as_float(__byte_perm(p_ev, magic, 0x7044)), g_acc[0]); g_acc[1] = fmaf(xv.x, __uint_as_float(__byte_perm(p_ev, magic, 0x7144)), g_acc[1]); g_acc[2] = fmaf(xv.x, __uint_as_float(__byte_perm(p_ev, magic, 0x7244)), g_acc[2]); g_acc[3] = fmaf(xv.x, __uint_as_float(__byte_perm(p_ev, magic, 0x7344)), g_acc[3]); g_acc[0] = fmaf(xv.y, __uint_as_float(__byte_perm(p_od, magic, 0x7044)), g_acc[0]); g_acc[1] = fmaf(xv.y, __uint_as_float(__byte_perm(p_od, magic, 0x7144)), g_acc[1]); g_acc[2] = fmaf(xv.y, __uint_as_float(__byte_perm(p_od, magic, 0x7244)), g_acc[2]); g_acc[3] = fmaf(xv.y, __uint_as_float(__byte_perm(p_od, magic, 0x7344)), g_acc[3]); } // Full group (8 kp-rows per warp, 64 kp-rows across 8 warps, 128 contiguous output columns per warp) __device__ __forceinline__ void w4a16_group_tile_128( const uint8_t* __restrict__ w_q, const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ zeros, const float* __restrict__ smem_x, int N, int n0, int g, int warp_id, int lane, float (&acc)[4] ) { const int kp_base = g * 64 + warp_id * 8; const int col = n0 + lane * 4; const uint8_t* w_ptr = w_q + col; const float2* x_f2 = reinterpret_cast(smem_x + kp_base * 2); float g_acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; float x_sum = 0.0f; #pragma unroll for (int r = 0; r < 8; ++r) { uint32_t p = *reinterpret_cast(w_ptr + (kp_base + r) * N); float2 xv = x_f2[r]; x_sum += xv.x + xv.y; uint32_t magic = 0x43000000u + (threadIdx.x >> 20); accum_u32_fast(p, xv, magic, g_acc); } const int sz_idx = g * N + col; uint2 s_u2 = *reinterpret_cast(scales + sz_idx); uint2 z_u2 = *reinterpret_cast(zeros + sz_idx); const __nv_bfloat16* s_bf = reinterpret_cast(&s_u2); const __nv_bfloat16* z_bf = reinterpret_cast(&z_u2); #pragma unroll for (int c = 0; c < 4; ++c) { float sv = bf16_to_f32(s_bf[c]); float zv = bf16_to_f32(z_bf[c]) + 128.0f; acc[c] += sv * (g_acc[c] - zv * x_sum); } } // Fused Gate+Up full-group tile (reads x_f2 and computes x_sum ONCE for both Gate and Up!) __device__ __forceinline__ void w4a16_gate_up_group_tile_128( const uint8_t* __restrict__ gw_q, const __nv_bfloat16* __restrict__ g_scales, const __nv_bfloat16* __restrict__ g_zeros, const uint8_t* __restrict__ uw_q, const __nv_bfloat16* __restrict__ u_scales, const __nv_bfloat16* __restrict__ u_zeros, const float* __restrict__ smem_x, int N, int n0, int g, int warp_id, int lane, float (&acc_g)[4], float (&acc_u)[4] ) { const int kp_base = g * 64 + warp_id * 8; const int col = n0 + lane * 4; const uint8_t* gw_ptr = gw_q + col; const uint8_t* uw_ptr = uw_q + col; const float2* x_f2 = reinterpret_cast(smem_x + kp_base * 2); float gg_acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; float ug_acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; float x_sum = 0.0f; #pragma unroll for (int r = 0; r < 8; ++r) { int row_off = (kp_base + r) * N; uint32_t gp = *reinterpret_cast(gw_ptr + row_off); uint32_t up = *reinterpret_cast(uw_ptr + row_off); float2 xv = x_f2[r]; x_sum += xv.x + xv.y; uint32_t magic = 0x43000000u + (threadIdx.x >> 20); accum_u32_fast(gp, xv, magic, gg_acc); accum_u32_fast(up, xv, magic, ug_acc); } const int sz_idx = g * N + col; uint2 gs_u2 = *reinterpret_cast(g_scales + sz_idx); uint2 gz_u2 = *reinterpret_cast(g_zeros + sz_idx); uint2 us_u2 = *reinterpret_cast(u_scales + sz_idx); uint2 uz_u2 = *reinterpret_cast(u_zeros + sz_idx); const __nv_bfloat16* gs_bf = reinterpret_cast(&gs_u2); const __nv_bfloat16* gz_bf = reinterpret_cast(&gz_u2); const __nv_bfloat16* us_bf = reinterpret_cast(&us_u2); const __nv_bfloat16* uz_bf = reinterpret_cast(&uz_u2); #pragma unroll for (int c = 0; c < 4; ++c) { acc_g[c] += bf16_to_f32(gs_bf[c]) * (gg_acc[c] - (bf16_to_f32(gz_bf[c]) + 128.0f) * x_sum); acc_u[c] += bf16_to_f32(us_bf[c]) * (ug_acc[c] - (bf16_to_f32(uz_bf[c]) + 128.0f) * x_sum); } } __device__ __forceinline__ void save_tile_128( float (&acc)[4], float* smem_red_prev, // [8 * 128] int warp_id, int lane ) { *reinterpret_cast(&smem_red_prev[warp_id * 128 + lane * 4]) = make_float4(acc[0], acc[1], acc[2], acc[3]); acc[0] = 0.0f; acc[1] = 0.0f; acc[2] = 0.0f; acc[3] = 0.0f; } __device__ __forceinline__ void flush_dual_tile_128( float* dst_cur, float* dst_prev, float (&acc)[4], float* smem_red_cur, // [8 * 128] float* smem_red_prev, // [8 * 128] int tid, int warp_id, int lane, float scale_cur = 1.0f, float scale_prev = 1.0f ) { *reinterpret_cast(&smem_red_cur[warp_id * 128 + lane * 4]) = make_float4(acc[0], acc[1], acc[2], acc[3]); __syncthreads(); if (tid < 128) { if (dst_cur != nullptr) { float sum = smem_red_cur[0 * 128 + tid] + smem_red_cur[1 * 128 + tid] + smem_red_cur[2 * 128 + tid] + smem_red_cur[3 * 128 + tid] + smem_red_cur[4 * 128 + tid] + smem_red_cur[5 * 128 + tid] + smem_red_cur[6 * 128 + tid] + smem_red_cur[7 * 128 + tid]; atomicAdd(&dst_cur[tid], sum * scale_cur); } } else if (dst_prev != nullptr) { int c = tid - 128; float sum = smem_red_prev[0 * 128 + c] + smem_red_prev[1 * 128 + c] + smem_red_prev[2 * 128 + c] + smem_red_prev[3 * 128 + c] + smem_red_prev[4 * 128 + c] + smem_red_prev[5 * 128 + c] + smem_red_prev[6 * 128 + c] + smem_red_prev[7 * 128 + c]; atomicAdd(&dst_prev[c], sum * scale_prev); } } __device__ __forceinline__ void save_tile_128_pair( float (&acc_g)[4], float (&acc_u)[4], float* smem_red_g_prev, float* smem_red_u_prev, int warp_id, int lane ) { *reinterpret_cast(&smem_red_g_prev[warp_id * 128 + lane * 4]) = make_float4(acc_g[0], acc_g[1], acc_g[2], acc_g[3]); *reinterpret_cast(&smem_red_u_prev[warp_id * 128 + lane * 4]) = make_float4(acc_u[0], acc_u[1], acc_u[2], acc_u[3]); acc_g[0] = 0.0f; acc_g[1] = 0.0f; acc_g[2] = 0.0f; acc_g[3] = 0.0f; acc_u[0] = 0.0f; acc_u[1] = 0.0f; acc_u[2] = 0.0f; acc_u[3] = 0.0f; } __device__ __forceinline__ void flush_dual_tile_128_pair( float* dst_g_cur, float* dst_u_cur, float* dst_g_prev, float* dst_u_prev, float (&acc_g)[4], float (&acc_u)[4], float* smem_red_g_cur, float* smem_red_u_cur, float* smem_red_g_prev, float* smem_red_u_prev, int tid, int warp_id, int lane ) { *reinterpret_cast(&smem_red_g_cur[warp_id * 128 + lane * 4]) = make_float4(acc_g[0], acc_g[1], acc_g[2], acc_g[3]); *reinterpret_cast(&smem_red_u_cur[warp_id * 128 + lane * 4]) = make_float4(acc_u[0], acc_u[1], acc_u[2], acc_u[3]); __syncthreads(); if (tid < 128) { if (dst_g_cur != nullptr) { float sum_g = smem_red_g_cur[0 * 128 + tid] + smem_red_g_cur[1 * 128 + tid] + smem_red_g_cur[2 * 128 + tid] + smem_red_g_cur[3 * 128 + tid] + smem_red_g_cur[4 * 128 + tid] + smem_red_g_cur[5 * 128 + tid] + smem_red_g_cur[6 * 128 + tid] + smem_red_g_cur[7 * 128 + tid]; atomicAdd(&dst_g_cur[tid], sum_g); } if (dst_g_prev != nullptr) { float sum_gp = smem_red_g_prev[0 * 128 + tid] + smem_red_g_prev[1 * 128 + tid] + smem_red_g_prev[2 * 128 + tid] + smem_red_g_prev[3 * 128 + tid] + smem_red_g_prev[4 * 128 + tid] + smem_red_g_prev[5 * 128 + tid] + smem_red_g_prev[6 * 128 + tid] + smem_red_g_prev[7 * 128 + tid]; atomicAdd(&dst_g_prev[tid], sum_gp); } } else { int c = tid - 128; if (dst_u_cur != nullptr) { float sum_u = smem_red_u_cur[0 * 128 + c] + smem_red_u_cur[1 * 128 + c] + smem_red_u_cur[2 * 128 + c] + smem_red_u_cur[3 * 128 + c] + smem_red_u_cur[4 * 128 + c] + smem_red_u_cur[5 * 128 + c] + smem_red_u_cur[6 * 128 + c] + smem_red_u_cur[7 * 128 + c]; atomicAdd(&dst_u_cur[c], sum_u); } if (dst_u_prev != nullptr) { float sum_up = smem_red_u_prev[0 * 128 + c] + smem_red_u_prev[1 * 128 + c] + smem_red_u_prev[2 * 128 + c] + smem_red_u_prev[3 * 128 + c] + smem_red_u_prev[4 * 128 + c] + smem_red_u_prev[5 * 128 + c] + smem_red_u_prev[6 * 128 + c] + smem_red_u_prev[7 * 128 + c]; atomicAdd(&dst_u_prev[c], sum_up); } } } __device__ __forceinline__ void flush_tile_128( float* dst_fp32, float (&acc)[4], float* smem_red, // [8 * 128] int tid, int warp_id, int lane, float scale_factor = 1.0f ) { *reinterpret_cast(&smem_red[warp_id * 128 + lane * 4]) = make_float4(acc[0], acc[1], acc[2], acc[3]); __syncthreads(); if (tid < 128) { float sum = smem_red[0 * 128 + tid] + smem_red[1 * 128 + tid] + smem_red[2 * 128 + tid] + smem_red[3 * 128 + tid] + smem_red[4 * 128 + tid] + smem_red[5 * 128 + tid] + smem_red[6 * 128 + tid] + smem_red[7 * 128 + tid]; atomicAdd(&dst_fp32[tid], sum * scale_factor); } __syncthreads(); acc[0] = 0.0f; acc[1] = 0.0f; acc[2] = 0.0f; acc[3] = 0.0f; } __device__ __forceinline__ void compute_moe_routing_smem( const float* __restrict__ router_logits, int* smem_top_idx, float* smem_top_w, int tid ) { if (tid < 32) { float l0 = router_logits[tid]; float l1 = router_logits[tid + 32]; float m = fmaxf(l0, l1); #pragma unroll for (int off = 16; off > 0; off >>= 1) { m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off)); } float p0 = expf(l0 - m); float p1 = expf(l1 - m); float s = p0 + p1; s = warp_reduce_sum(s); p0 /= s; p1 /= s; float w_sum = 0.0f; #pragma unroll for (int j = 0; j < MOE_ACTIVE; ++j) { float best_v = (p0 >= p1) ? p0 : p1; int best_i = (p0 >= p1) ? tid : (tid + 32); #pragma unroll for (int off = 16; off > 0; off >>= 1) { float other_v = __shfl_xor_sync(0xffffffff, best_v, off); int other_i = __shfl_xor_sync(0xffffffff, best_i, off); if (other_v > best_v || (other_v == best_v && other_i < best_i)) { best_v = other_v; best_i = other_i; } } if (tid == 0) { smem_top_idx[j] = best_i; smem_top_w[j] = best_v; } w_sum += best_v; if (best_i == tid) p0 = -1.0f; if (best_i == tid + 32) p1 = -1.0f; } if (tid < MOE_ACTIVE) { smem_top_w[tid] = smem_top_w[tid] / (w_sum + 1.0e-9f) * 2.446f; } if (tid == 0) { smem_top_idx[MOE_ACTIVE] = 0; smem_top_w[MOE_ACTIVE] = 1.0f; } } } __device__ void run_moe_phases( uint32_t& local_epoch, const __nv_bfloat16* __restrict__ moe_norm, const MoEPtrs& moe, const ScratchPtrs& sc, float* smem_x, // [6144] floats float* smem_red, // [1024] floats int* smem_top_idx, float* smem_top_w ) { const int bid = blockIdx.x; const int tid = threadIdx.x; const int warp_id = tid >> 5; const int lane = tid & 31; // ---------------- Phase E1: Attn Residual + MoE RMSNorm + Router + Zeroing ---------------- if (bid < 64) { const __nv_bfloat16* __restrict__ res_in = sc.res_h + HIDDEN; const float* __restrict__ a_in = sc.attn_out_fp32; const __nv_bfloat16* __restrict__ rw_row = moe.router_w + bid * HIDDEN; float hv[9]; __nv_bfloat16 h_bf[9]; float m_norm[9]; float rw_val[9]; float local_sq = 0.0f; #pragma unroll for (int m = 0; m < 9; ++m) { int i = tid + m * BLOCK_THREADS; float h_old = bf16_to_f32(res_in[i]); float a_out = round_bf16(a_in[i]); h_bf[m] = f32_to_bf16(h_old + a_out); m_norm[m] = bf16_to_f32(moe_norm[i]); rw_val[m] = bf16_to_f32(rw_row[i]); hv[m] = bf16_to_f32(h_bf[m]); local_sq += hv[m] * hv[m]; } if (bid == 0) { #pragma unroll for (int m = 0; m < 9; ++m) { sc.res_h[tid + m * BLOCK_THREADS] = h_bf[m]; } } local_sq = warp_reduce_sum(local_sq); if (lane == 0) smem_red[warp_id] = local_sq; __syncthreads(); if (tid < 32) { float b_sq = (tid < 8) ? smem_red[tid] : 0.0f; b_sq = warp_reduce_sum(b_sq); if (tid == 0) smem_red[0] = rsqrtf(b_sq / float(HIDDEN) + 1.0e-6f); } __syncthreads(); float inv_rms = smem_red[0]; float r_dot = 0.0f; #pragma unroll for (int m = 0; m < 9; ++m) { int i = tid + m * BLOCK_THREADS; float xn = round_bf16(hv[m] * inv_rms * m_norm[m]); smem_x[i] = xn; if (bid == 0) { sc.norm_x[i] = xn; } r_dot += xn * rw_val[m]; } r_dot = warp_reduce_sum(r_dot); if (lane == 0) smem_red[warp_id] = r_dot; __syncthreads(); if (tid < 32) { float sum = (tid < 8) ? smem_red[tid] : 0.0f; sum = warp_reduce_sum(sum); if (tid == 0) sc.router_logits[bid] = round_bf16(sum); } } else { const int zero_tid = (bid - 64) * BLOCK_THREADS + tid; const int zero_stride = (NUM_SMS - 64) * BLOCK_THREADS; for (int i = zero_tid; i < MOE_TOTAL_ACTIVE * MOE_INTER; i += zero_stride) { sc.moe_gate_fp32[i] = 0.0f; sc.moe_up_fp32[i] = 0.0f; } for (int i = zero_tid; i < HIDDEN; i += zero_stride) { sc.moe_out_fp32[i] = 0.0f; } } grid_barrier(local_epoch); // ---------------- Phase E2: Top-8 Selection + Fused Gate & Up W4A16 GEMV ---------------- compute_moe_routing_smem(sc.router_logits, smem_top_idx, smem_top_w, tid); if (bid >= 64) { for (int i = tid; i < HIDDEN; i += BLOCK_THREADS) { smem_x[i] = sc.norm_x[i]; } } __syncthreads(); float* smem_red_u = smem_x + 2304; float* smem_red_g_prev = smem_x + 3328; float* smem_red_u_prev = smem_x + 4352; const int eff_bid = (bid < 188) ? (bid << 1) : (((bid - 188) << 1) + 1); // 9 active experts * 8 col_tiles (1024/128) * 18 full_groups (2304/128) = 1296 units constexpr int TOTAL_GU_GROUPS = MOE_TOTAL_ACTIVE * 8 * 18; const int gu_start = (eff_bid * TOTAL_GU_GROUPS) / NUM_SMS; const int gu_end = ((eff_bid + 1) * TOTAL_GU_GROUPS) / NUM_SMS; float acc_g[4] = {0.0f, 0.0f, 0.0f, 0.0f}; float acc_u[4] = {0.0f, 0.0f, 0.0f, 0.0f}; int cur_sc = -1; int prev_sc = -1; for (int u = gu_start; u < gu_end; ++u) { int sc_idx = u / 18; int g = u % 18; if (sc_idx != cur_sc) { if (cur_sc >= 0) { save_tile_128_pair(acc_g, acc_u, smem_red_g_prev, smem_red_u_prev, warp_id, lane); prev_sc = cur_sc; } cur_sc = sc_idx; } int s = sc_idx >> 3; int c_tile = sc_idx & 7; int e = smem_top_idx[s]; const QLinearPtrs& g_ptrs = (s < MOE_ACTIVE) ? moe.gate : moe.s_gate; const QLinearPtrs& u_ptrs = (s < MOE_ACTIVE) ? moe.up : moe.s_up; const size_t wq_off = size_t(e) * (HIDDEN / 2) * MOE_INTER; const size_t sz_off = size_t(e) * (HIDDEN / GROUP) * MOE_INTER; w4a16_gate_up_group_tile_128( g_ptrs.w_q + wq_off, g_ptrs.scales + sz_off, g_ptrs.zeros + sz_off, u_ptrs.w_q + wq_off, u_ptrs.scales + sz_off, u_ptrs.zeros + sz_off, smem_x, MOE_INTER, c_tile * 128, g, warp_id, lane, acc_g, acc_u ); } flush_dual_tile_128_pair( (cur_sc >= 0) ? (sc.moe_gate_fp32 + cur_sc * 128) : nullptr, (cur_sc >= 0) ? (sc.moe_up_fp32 + cur_sc * 128) : nullptr, (prev_sc >= 0) ? (sc.moe_gate_fp32 + prev_sc * 128) : nullptr, (prev_sc >= 0) ? (sc.moe_up_fp32 + prev_sc * 128) : nullptr, acc_g, acc_u, smem_red, smem_red_u, smem_red_g_prev, smem_red_u_prev, tid, warp_id, lane ); grid_barrier(local_epoch); // ---------------- Phase E3: MoE Down W4A16 GEMV (Per-Expert Full-Group Mapped, 41 blocks/expert) ---------------- if (eff_bid < 369) { const int s_my = eff_bid / 41; const int b_in_s = eff_bid % 41; const int u_start = (b_in_s * 144) / 41; const int u_end = ((b_in_s + 1) * 144) / 41; const float4 g4 = reinterpret_cast(sc.moe_gate_fp32 + s_my * MOE_INTER)[tid]; const float4 u4 = reinterpret_cast(sc.moe_up_fp32 + s_my * MOE_INTER)[tid]; reinterpret_cast(smem_x)[tid] = make_float4( (g4.x / (1.0f + expf(-g4.x))) * u4.x, (g4.y / (1.0f + expf(-g4.y))) * u4.y, (g4.z / (1.0f + expf(-g4.z))) * u4.z, (g4.w / (1.0f + expf(-g4.w))) * u4.w ); __syncthreads(); const int e = smem_top_idx[s_my]; const float w_scale = smem_top_w[s_my]; const QLinearPtrs& d_ptrs = (s_my < MOE_ACTIVE) ? moe.down : moe.s_down; const uint8_t* dw_q = d_ptrs.w_q + size_t(e) * (MOE_INTER / 2) * HIDDEN; const __nv_bfloat16* d_scales = d_ptrs.scales + size_t(e) * (MOE_INTER / GROUP) * HIDDEN; const __nv_bfloat16* d_zeros = d_ptrs.zeros + size_t(e) * (MOE_INTER / GROUP) * HIDDEN; float acc_d[4] = {0.0f, 0.0f, 0.0f, 0.0f}; int cur_ct = -1; int prev_ct = -1; for (int u = u_start; u < u_end; ++u) { int c_tile = u >> 3; int g = u & 7; if (c_tile != cur_ct) { if (cur_ct >= 0) { save_tile_128(acc_d, smem_red_g_prev, warp_id, lane); prev_ct = cur_ct; } cur_ct = c_tile; } w4a16_group_tile_128( dw_q, d_scales, d_zeros, smem_x, HIDDEN, c_tile * 128, g, warp_id, lane, acc_d ); } flush_dual_tile_128( (cur_ct >= 0) ? (sc.moe_out_fp32 + cur_ct * 128) : nullptr, (prev_ct >= 0) ? (sc.moe_out_fp32 + prev_ct * 128) : nullptr, acc_d, smem_red, smem_red_g_prev, tid, warp_id, lane, w_scale, w_scale ); } grid_barrier(local_epoch); } __global__ __launch_bounds__(BLOCK_THREADS, 2) void kimi_decode_megakernel(KernelParams params) { uint32_t local_epoch = (threadIdx.x == 0) ? *reinterpret_cast(&g_sync_arr[1]) : 0u; const int bid = blockIdx.x; const int tid = threadIdx.x; const int warp_id = tid >> 5; const int lane = tid & 31; const int eff_bid = (bid < 188) ? (bid << 1) : (((bid - 188) << 1) + 1); __shared__ float smem_x[6144]; __shared__ float smem_red[8 * 128]; __shared__ int smem_top_idx[16]; __shared__ float smem_top_w[16]; ScratchPtrs sc = params.scratch; float* smem_red_prev = smem_x + 4096; // ========================================================================= // BLOCKS 0, 1, 2: KDA + MoE // ========================================================================= for (int b = 0; b < 3; ++b) { const KDABlockPtrs kb = params.kda[b]; // ---------------- Phase K1: Residual Add + Pre-Attn RMSNorm + Beta_proj (blocks 0..31) + Zeroing ---------------- if (bid < 32) { const __nv_bfloat16* __restrict__ res_in = sc.res_h; const float* __restrict__ m_in = sc.moe_out_fp32; const __nv_bfloat16* __restrict__ a_norm = kb.attn_norm; const __nv_bfloat16* __restrict__ bw = kb.beta_w + bid * HIDDEN; float hv[9], an_val[9], bw_val[9]; __nv_bfloat16 h_bf[9]; float local_sq = 0.0f; #pragma unroll for (int m = 0; m < 9; ++m) { int i = tid + m * BLOCK_THREADS; if (b == 0) { h_bf[m] = params.hidden_in[i]; } else { float h_old = bf16_to_f32(res_in[i]); float m_out = round_bf16(m_in[i]); h_bf[m] = f32_to_bf16(h_old + m_out); } an_val[m] = bf16_to_f32(a_norm[i]); bw_val[m] = bf16_to_f32(bw[i]); hv[m] = bf16_to_f32(h_bf[m]); local_sq += hv[m] * hv[m]; } if (bid == 0) { #pragma unroll for (int m = 0; m < 9; ++m) { (sc.res_h + HIDDEN)[tid + m * BLOCK_THREADS] = h_bf[m]; } } local_sq = warp_reduce_sum(local_sq); if (lane == 0) smem_red[warp_id] = local_sq; __syncthreads(); if (tid < 32) { float b_sq = (tid < 8) ? smem_red[tid] : 0.0f; b_sq = warp_reduce_sum(b_sq); if (tid == 0) smem_red[0] = rsqrtf(b_sq / float(HIDDEN) + 1.0e-6f); } __syncthreads(); float inv_rms = smem_red[0]; float b_dot = 0.0f; #pragma unroll for (int m = 0; m < 9; ++m) { int i = tid + m * BLOCK_THREADS; float xn = round_bf16(hv[m] * inv_rms * an_val[m]); smem_x[i] = xn; if (bid == 0) { sc.norm_x[i] = xn; } b_dot += xn * bw_val[m]; } b_dot = warp_reduce_sum(b_dot); if (lane == 0) smem_red[warp_id] = b_dot; __syncthreads(); if (tid < 32) { float sum = (tid < 8) ? smem_red[tid] : 0.0f; sum = warp_reduce_sum(sum); if (tid == 0) { float b_bf = round_bf16(sum); sc.kda_beta[bid] = 1.0f / (1.0f + expf(-b_bf)); } } } else { const int z_tid = (bid - 32) * BLOCK_THREADS + tid; const int z_stride = (NUM_SMS - 32) * BLOCK_THREADS; for (int i = z_tid; i < 4 * KDA_C; i += z_stride) { sc.kda_qkvg_fp32[i] = 0.0f; } for (int i = z_tid; i < HIDDEN; i += z_stride) { sc.attn_out_fp32[i] = 0.0f; } } grid_barrier(local_epoch); // ---------------- Phase K2: KDA q/k/v/g W4A16 GEMV ---------------- if (bid >= 32) { for (int i = tid; i < HIDDEN; i += BLOCK_THREADS) { smem_x[i] = sc.norm_x[i]; } __syncthreads(); } constexpr int TOTAL_QKVG_UNITS = 4 * 32 * 18; const int qk_start = (eff_bid * TOTAL_QKVG_UNITS) / NUM_SMS; const int qk_end = ((eff_bid + 1) * TOTAL_QKVG_UNITS) / NUM_SMS; float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; int cur_pc = -1; int prev_pc = -1; for (int u = qk_start; u < qk_end; ++u) { int pc_idx = u / 18; int g = u % 18; if (pc_idx != cur_pc) { if (cur_pc >= 0) { save_tile_128(acc, smem_red_prev, warp_id, lane); prev_pc = cur_pc; } cur_pc = pc_idx; } int p_id = pc_idx >> 5; int c_tile = pc_idx & 31; const QLinearPtrs& ql = (p_id == 0) ? kb.q_proj : (p_id == 1) ? kb.k_proj : (p_id == 2) ? kb.v_proj : kb.g_proj; w4a16_group_tile_128( ql.w_q, ql.scales, ql.zeros, smem_x, KDA_C, c_tile * 128, g, warp_id, lane, acc ); } flush_dual_tile_128( (cur_pc >= 0) ? (sc.kda_qkvg_fp32 + cur_pc * 128) : nullptr, (prev_pc >= 0) ? (sc.kda_qkvg_fp32 + prev_pc * 128) : nullptr, acc, smem_red, smem_red_prev, tid, warp_id, lane ); grid_barrier(local_epoch); // ---------------- Phase K3: Short Causal Conv + KDA Recurrent State S Update (128 blocks) ---------------- if (bid < 128) { const int h = bid >> 2; const int bj = bid & 3; float* s_q = smem_x + 0; float* s_k = smem_x + 128; float* s_v = smem_x + 256; float* s_g = smem_x + 384; if (tid < 128) { const int c = h * 128 + tid; float q_raw = round_bf16(sc.kda_qkvg_fp32[0 * KDA_C + c]); const __nv_bfloat16* cw_q = kb.conv_w + (0 * KDA_C + c) * 4; float q_c = bf16_to_f32(kb.cq[0 * KDA_C + c]) * bf16_to_f32(cw_q[0]) + bf16_to_f32(kb.cq[1 * KDA_C + c]) * bf16_to_f32(cw_q[1]) + bf16_to_f32(kb.cq[2 * KDA_C + c]) * bf16_to_f32(cw_q[2]) + q_raw * bf16_to_f32(cw_q[3]); float q_silu = round_bf16(q_c / (1.0f + expf(-q_c))); s_q[tid] = q_silu * 0.08838834764831845f; float k_raw = round_bf16(sc.kda_qkvg_fp32[1 * KDA_C + c]); const __nv_bfloat16* cw_k = kb.conv_w + (1 * KDA_C + c) * 4; float k_c = bf16_to_f32(kb.ck[0 * KDA_C + c]) * bf16_to_f32(cw_k[0]) + bf16_to_f32(kb.ck[1 * KDA_C + c]) * bf16_to_f32(cw_k[1]) + bf16_to_f32(kb.ck[2 * KDA_C + c]) * bf16_to_f32(cw_k[2]) + k_raw * bf16_to_f32(cw_k[3]); s_k[tid] = round_bf16(k_c / (1.0f + expf(-k_c))); float v_raw = round_bf16(sc.kda_qkvg_fp32[2 * KDA_C + c]); const __nv_bfloat16* cw_v = kb.conv_w + (2 * KDA_C + c) * 4; float v_c = bf16_to_f32(kb.cv[0 * KDA_C + c]) * bf16_to_f32(cw_v[0]) + bf16_to_f32(kb.cv[1 * KDA_C + c]) * bf16_to_f32(cw_v[1]) + bf16_to_f32(kb.cv[2 * KDA_C + c]) * bf16_to_f32(cw_v[2]) + v_raw * bf16_to_f32(cw_v[3]); s_v[tid] = round_bf16(v_c / (1.0f + expf(-v_c))); float g_raw = round_bf16(sc.kda_qkvg_fp32[3 * KDA_C + c]); float sp = (g_raw > 20.0f) ? g_raw : log1pf(expf(g_raw)); s_g[tid] = expf(-sp); } __syncthreads(); const int tid_i = tid >> 5; const int tid_j = tid & 31; const int j = bj * 32 + tid_j; float* S_head = kb.S + h * (128 * 128); float s_reg[16]; float pred_part = 0.0f; #pragma unroll for (int m = 0; m < 16; ++m) { int i = tid_i * 16 + m; float sv = S_head[i * 128 + j] * s_g[i]; s_reg[m] = sv; pred_part += sv * s_k[i]; } smem_red[tid_i * 32 + tid_j] = pred_part; __syncthreads(); float pred = smem_red[0 * 32 + tid_j] + smem_red[1 * 32 + tid_j] + smem_red[2 * 32 + tid_j] + smem_red[3 * 32 + tid_j] + smem_red[4 * 32 + tid_j] + smem_red[5 * 32 + tid_j] + smem_red[6 * 32 + tid_j] + smem_red[7 * 32 + tid_j]; float diff_j = sc.kda_beta[h] * (s_v[j] - pred); float o_part = 0.0f; #pragma unroll for (int m = 0; m < 16; ++m) { int i = tid_i * 16 + m; float sv = s_reg[m] + s_k[i] * diff_j; S_head[i * 128 + j] = sv; o_part += sv * s_q[i]; } smem_red[tid_i * 32 + tid_j] = o_part; __syncthreads(); if (tid_i == 0) { float o_val = smem_red[0 * 32 + tid_j] + smem_red[1 * 32 + tid_j] + smem_red[2 * 32 + tid_j] + smem_red[3 * 32 + tid_j] + smem_red[4 * 32 + tid_j] + smem_red[5 * 32 + tid_j] + smem_red[6 * 32 + tid_j] + smem_red[7 * 32 + tid_j]; sc.kda_o[h * 128 + j] = round_bf16(o_val); } } grid_barrier(local_epoch); // ---------------- Phase K4: KDA o_proj W4A16 GEMV + Conv Window Update ---------------- if (bid < 16) { int c = bid * BLOCK_THREADS + tid; kb.cq[0 * KDA_C + c] = kb.cq[1 * KDA_C + c]; kb.cq[1 * KDA_C + c] = kb.cq[2 * KDA_C + c]; kb.cq[2 * KDA_C + c] = f32_to_bf16(sc.kda_qkvg_fp32[0 * KDA_C + c]); kb.ck[0 * KDA_C + c] = kb.ck[1 * KDA_C + c]; kb.ck[1 * KDA_C + c] = kb.ck[2 * KDA_C + c]; kb.ck[2 * KDA_C + c] = f32_to_bf16(sc.kda_qkvg_fp32[1 * KDA_C + c]); kb.cv[0 * KDA_C + c] = kb.cv[1 * KDA_C + c]; kb.cv[1 * KDA_C + c] = kb.cv[2 * KDA_C + c]; kb.cv[2 * KDA_C + c] = f32_to_bf16(sc.kda_qkvg_fp32[2 * KDA_C + c]); } for (int i = tid; i < KDA_C; i += BLOCK_THREADS) { smem_x[i] = sc.kda_o[i]; } __syncthreads(); constexpr int TOTAL_O_GROUPS = 18 * 32; const int o_start = (eff_bid * TOTAL_O_GROUPS) / NUM_SMS; const int o_end = ((eff_bid + 1) * TOTAL_O_GROUPS) / NUM_SMS; float acc_o[4] = {0.0f, 0.0f, 0.0f, 0.0f}; int cur_oc = -1; int prev_oc = -1; for (int u = o_start; u < o_end; ++u) { int c_tile = u >> 5; int g = u & 31; if (c_tile != cur_oc) { if (cur_oc >= 0) { save_tile_128(acc_o, smem_red_prev, warp_id, lane); prev_oc = cur_oc; } cur_oc = c_tile; } w4a16_group_tile_128( kb.o_proj.w_q, kb.o_proj.scales, kb.o_proj.zeros, smem_x, HIDDEN, c_tile * 128, g, warp_id, lane, acc_o ); } flush_dual_tile_128( (cur_oc >= 0) ? (sc.attn_out_fp32 + cur_oc * 128) : nullptr, (prev_oc >= 0) ? (sc.attn_out_fp32 + prev_oc * 128) : nullptr, acc_o, smem_red, smem_red_prev, tid, warp_id, lane ); grid_barrier(local_epoch); run_moe_phases(local_epoch, kb.moe_norm, kb.moe, sc, smem_x, smem_red, smem_top_idx, smem_top_w); } // ========================================================================= // BLOCK 3: MLA + MoE // ========================================================================= const MLABlockPtrs mb = params.mla; // ---------------- Phase M1: Residual Add + Pre-Attn RMSNorm + Cache Copy ---------------- if (bid == 0) { const __nv_bfloat16* __restrict__ res_in = sc.res_h; const float* __restrict__ m_in = sc.moe_out_fp32; const __nv_bfloat16* __restrict__ a_norm = mb.attn_norm; float hv[9], an_val[9]; __nv_bfloat16 h_bf[9]; float local_sq = 0.0f; #pragma unroll for (int m = 0; m < 9; ++m) { int i = tid + m * BLOCK_THREADS; float h_old = bf16_to_f32(res_in[i]); float m_out = round_bf16(m_in[i]); h_bf[m] = f32_to_bf16(h_old + m_out); an_val[m] = bf16_to_f32(a_norm[i]); hv[m] = bf16_to_f32(h_bf[m]); local_sq += hv[m] * hv[m]; } #pragma unroll for (int m = 0; m < 9; ++m) { (sc.res_h + HIDDEN)[tid + m * BLOCK_THREADS] = h_bf[m]; } local_sq = warp_reduce_sum(local_sq); if (lane == 0) smem_red[warp_id] = local_sq; __syncthreads(); if (tid < 32) { float b_sq = (tid < 8) ? smem_red[tid] : 0.0f; b_sq = warp_reduce_sum(b_sq); if (tid == 0) smem_red[0] = rsqrtf(b_sq / float(HIDDEN) + 1.0e-6f); } __syncthreads(); float inv_rms = smem_red[0]; #pragma unroll for (int m = 0; m < 9; ++m) { int i = tid + m * BLOCK_THREADS; float xn = round_bf16(hv[m] * inv_rms * an_val[m]); smem_x[i] = xn; sc.norm_x[i] = xn; } } else { const int z_tid = (bid - 1) * BLOCK_THREADS + tid; const int z_stride = (NUM_SMS - 1) * BLOCK_THREADS; for (int i = z_tid; i < 53 * 128; i += z_stride) { sc.mla_q_kv_fp32[i] = 0.0f; } for (int i = z_tid; i < HIDDEN; i += z_stride) { sc.attn_out_fp32[i] = 0.0f; } if (mb.c_kv_src != mb.c_kv_dst) { const size_t total_ckv_u4 = (size_t(mb.pos) * KV_LORA) / 8; const uint4* src_ckv = reinterpret_cast(mb.c_kv_src); uint4* dst_ckv = reinterpret_cast(mb.c_kv_dst); for (size_t i = z_tid; i < total_ckv_u4; i += z_stride) { dst_ckv[i] = src_ckv[i]; } const size_t total_kr_u4 = (size_t(mb.pos) * QK_ROPE) / 8; const uint4* src_kr = reinterpret_cast(mb.k_rope_src); uint4* dst_kr = reinterpret_cast(mb.k_rope_dst); for (size_t i = z_tid; i < total_kr_u4; i += z_stride) { dst_kr[i] = src_kr[i]; } } } grid_barrier(local_epoch); // ---------------- Phase M2: MLA q_proj + kv_a W4A16 GEMV ---------------- if (bid > 0) { for (int i = tid; i < HIDDEN; i += BLOCK_THREADS) { smem_x[i] = sc.norm_x[i]; } __syncthreads(); } { constexpr int TOTAL_QKV_GROUPS = 53 * 18; const int mq_start = (eff_bid * TOTAL_QKV_GROUPS) / NUM_SMS; const int mq_end = ((eff_bid + 1) * TOTAL_QKV_GROUPS) / NUM_SMS; float acc_mq[4] = {0.0f, 0.0f, 0.0f, 0.0f}; int cur_mqc = -1; int prev_mqc = -1; for (int u = mq_start; u < mq_end; ++u) { int c_tile = u / 18; int g = u % 18; if (c_tile != cur_mqc) { if (cur_mqc >= 0) { save_tile_128(acc_mq, smem_red_prev, warp_id, lane); prev_mqc = cur_mqc; } cur_mqc = c_tile; } if (c_tile < 48) { w4a16_group_tile_128( mb.q_proj.w_q, mb.q_proj.scales, mb.q_proj.zeros, smem_x, 6144, c_tile * 128, g, warp_id, lane, acc_mq ); } else { int eff_lane = (c_tile < 52) ? lane : (lane & 15); w4a16_group_tile_128( mb.kv_a.w_q, mb.kv_a.scales, mb.kv_a.zeros, smem_x, 576, (c_tile - 48) * 128, g, warp_id, eff_lane, acc_mq ); } } flush_dual_tile_128( (cur_mqc >= 0) ? (sc.mla_q_kv_fp32 + cur_mqc * 128) : nullptr, (prev_mqc >= 0) ? (sc.mla_q_kv_fp32 + prev_mqc * 128) : nullptr, acc_mq, smem_red, smem_red_prev, tid, warp_id, lane ); } grid_barrier(local_epoch); // ---------------- Phase M3: Absorbed Query + RoPE + KV Cache Append + Zeroing kda_o (256 blocks) ---------------- if (bid < 256) { const int h = bid >> 3; const int g_half = bid & 7; const int g = g_half >> 1; float* s_q_scaled = smem_x; if (tid < 128) { float qn = round_bf16(sc.mla_q_kv_fp32[h * Q_DIM + tid]); int col = h * KVB_DIM + tid; float sv = bf16_to_f32(mb.kv_b.scales[g * 8192 + col]); float zv = bf16_to_f32(mb.kv_b.zeros[g * 8192 + col]); float qs = qn * sv; s_q_scaled[tid] = qs; smem_red[tid] = qs * zv; } __syncthreads(); if (tid < 32) { float zsum = smem_red[tid] + smem_red[tid + 32] + smem_red[tid + 64] + smem_red[tid + 96]; zsum = warp_reduce_sum(zsum); if (tid == 0) smem_red[0] = zsum; } __syncthreads(); const float q_zero_sum = smem_red[0]; const float4 qs4 = *reinterpret_cast(&s_q_scaled[lane * 4]); const int col_base = h * KVB_DIM + lane * 4; #pragma unroll for (int m = 0; m < 4; ++m) { int rp = g_half * 32 + warp_id * 4 + m; uint32_t p = *reinterpret_cast(&mb.kv_b.w_q[rp * 8192 + col_base]); float dot_ev = qs4.x * float(p & 0x0Fu) + qs4.y * float((p >> 8) & 0x0Fu) + qs4.z * float((p >> 16) & 0x0Fu) + qs4.w * float((p >> 24) & 0x0Fu); float dot_od = qs4.x * float((p >> 4) & 0x0Fu) + qs4.y * float((p >> 12) & 0x0Fu) + qs4.z * float((p >> 20) & 0x0Fu) + qs4.w * float((p >> 28) & 0x0Fu); dot_ev = warp_reduce_sum(dot_ev); dot_od = warp_reduce_sum(dot_od); if (lane == 0) { sc.mla_q_abs[h * KV_LORA + rp * 2 + 0] = dot_ev - q_zero_sum; sc.mla_q_abs[h * KV_LORA + rp * 2 + 1] = dot_od - q_zero_sum; } } } else if (bid == 256) { for (int idx = tid; idx < MLA_H * (QK_ROPE / 2); idx += BLOCK_THREADS) { int h = idx >> 5; int pair = idx & 31; float inv_freq = powf(10000.0f, -float(2 * pair) / float(QK_ROPE)); float ang = float(mb.pos) * inv_freq; float c_v = cosf(ang); float s_v = sinf(ang); int base = h * Q_DIM + QK_NOPE + 2 * pair; float ev = round_bf16(sc.mla_q_kv_fp32[base + 0]); float od = round_bf16(sc.mla_q_kv_fp32[base + 1]); sc.mla_q_rope[h * QK_ROPE + 2 * pair + 0] = round_bf16(ev * c_v - od * s_v); sc.mla_q_rope[h * QK_ROPE + 2 * pair + 1] = round_bf16(od * c_v + ev * s_v); } } else if (bid == 257) { __nv_bfloat16* dst_ckv = mb.c_kv_dst + size_t(mb.pos) * KV_LORA; for (int r = tid; r < KV_LORA; r += BLOCK_THREADS) { dst_ckv[r] = f32_to_bf16(sc.mla_q_kv_fp32[6144 + r]); } if (tid < (QK_ROPE / 2)) { int pair = tid; float inv_freq = powf(10000.0f, -float(2 * pair) / float(QK_ROPE)); float ang = float(mb.pos) * inv_freq; float c_v = cosf(ang); float s_v = sinf(ang); int base = 6144 + KV_LORA + 2 * pair; float ev = round_bf16(sc.mla_q_kv_fp32[base + 0]); float od = round_bf16(sc.mla_q_kv_fp32[base + 1]); __nv_bfloat16* dst_kr = mb.k_rope_dst + size_t(mb.pos) * QK_ROPE; dst_kr[2 * pair + 0] = f32_to_bf16(ev * c_v - od * s_v); dst_kr[2 * pair + 1] = f32_to_bf16(od * c_v + ev * s_v); } } else { const int z_tid = (bid - 258) * BLOCK_THREADS + tid; const int z_stride = (NUM_SMS - 258) * BLOCK_THREADS; for (int i = z_tid; i < 4096; i += z_stride) { sc.kda_o[i] = 0.0f; } } grid_barrier(local_epoch); // ---------------- Phase M4: 188-Split 2-Head/Warp (376 blocks) Cooperative-Smem FlashDecoding ---------------- { const int split_id = bid >> 1; const int h_group = bid & 1; const int h0 = h_group * 16 + warp_id * 2 + 0; const int h1 = h_group * 16 + warp_id * 2 + 1; const int total_len = mb.pos + 1; const int l_start = (split_id * total_len) / NUM_L_SPLITS; const int l_end = ((split_id + 1) * total_len) / NUM_L_SPLITS; float q_abs_0[16], q_abs_1[16]; const float* qa0 = sc.mla_q_abs + h0 * KV_LORA + lane * 4; const float* qa1 = sc.mla_q_abs + h1 * KV_LORA + lane * 4; #pragma unroll for (int c = 0; c < 4; ++c) { float4 v0 = *reinterpret_cast(qa0 + c * 128); float4 v1 = *reinterpret_cast(qa1 + c * 128); q_abs_0[c * 4 + 0] = v0.x; q_abs_0[c * 4 + 1] = v0.y; q_abs_0[c * 4 + 2] = v0.z; q_abs_0[c * 4 + 3] = v0.w; q_abs_1[c * 4 + 0] = v1.x; q_abs_1[c * 4 + 1] = v1.y; q_abs_1[c * 4 + 2] = v1.z; q_abs_1[c * 4 + 3] = v1.w; } float2 qr0 = *reinterpret_cast(sc.mla_q_rope + h0 * QK_ROPE + lane * 2); float2 qr1 = *reinterpret_cast(sc.mla_q_rope + h1 * QK_ROPE + lane * 2); float o_acc_0[16], o_acc_1[16]; #pragma unroll for (int m = 0; m < 16; ++m) { o_acc_0[m] = 0.0f; o_acc_1[m] = 0.0f; } float m_i0 = -1.0e30f, d_i0 = 0.0f; float m_i1 = -1.0e30f, d_i1 = 0.0f; constexpr float mla_scale = 0.07216878364870322f; uint4* smem_ckv_u4 = reinterpret_cast(smem_x); uint32_t* smem_kr = reinterpret_cast(smem_ckv_u4 + 512); for (int l_base = l_start; l_base < l_end; l_base += 8) { int l_tok = l_base + warp_id; if (l_tok < l_end) { const uint4* g_ckv = reinterpret_cast(mb.c_kv_dst + size_t(l_tok) * KV_LORA); smem_ckv_u4[warp_id * 64 + lane] = g_ckv[lane]; smem_ckv_u4[warp_id * 64 + 32 + lane] = g_ckv[32 + lane]; const uint32_t* g_kr = reinterpret_cast(mb.k_rope_dst + size_t(l_tok) * QK_ROPE); smem_kr[warp_id * 32 + lane] = g_kr[lane]; } __syncthreads(); int rem = l_end - l_base; if (rem > 8) rem = 8; for (int w = 0; w < rem; ++w) { const __nv_bfloat16* s_ckv = reinterpret_cast(smem_ckv_u4 + w * 64) + lane * 4; uint32_t kr_u32 = smem_kr[w * 32 + lane]; const __nv_bfloat16* kr_bf = reinterpret_cast(&kr_u32); float kr0 = bf16_to_f32(kr_bf[0]); float kr1 = bf16_to_f32(kr_bf[1]); float c_val[16]; #pragma unroll for (int c = 0; c < 4; ++c) { uint2 u2 = *reinterpret_cast(s_ckv + c * 128); const __nv_bfloat16* bf = reinterpret_cast(&u2); c_val[c * 4 + 0] = bf16_to_f32(bf[0]); c_val[c * 4 + 1] = bf16_to_f32(bf[1]); c_val[c * 4 + 2] = bf16_to_f32(bf[2]); c_val[c * 4 + 3] = bf16_to_f32(bf[3]); } float dot0 = qr0.x * kr0 + qr0.y * kr1; float dot1 = qr1.x * kr0 + qr1.y * kr1; #pragma unroll for (int m = 0; m < 16; ++m) { dot0 += q_abs_0[m] * c_val[m]; dot1 += q_abs_1[m] * c_val[m]; } float s0 = warp_reduce_sum(dot0) * mla_scale; float s1 = warp_reduce_sum(dot1) * mla_scale; float mn0 = fmaxf(m_i0, s0); float al0 = expf(m_i0 - mn0); float p0 = expf(s0 - mn0); d_i0 = d_i0 * al0 + p0; m_i0 = mn0; float mn1 = fmaxf(m_i1, s1); float al1 = expf(m_i1 - mn1); float p1 = expf(s1 - mn1); d_i1 = d_i1 * al1 + p1; m_i1 = mn1; #pragma unroll for (int m = 0; m < 16; ++m) { o_acc_0[m] = o_acc_0[m] * al0 + p0 * c_val[m]; o_acc_1[m] = o_acc_1[m] * al1 + p1 * c_val[m]; } } __syncthreads(); } const int sh0 = split_id * MLA_H + h0; const int sh1 = split_id * MLA_H + h1; if (lane == 0) { sc.mla_att_m[sh0] = m_i0; sc.mla_att_d[sh0] = d_i0; sc.mla_att_m[sh1] = m_i1; sc.mla_att_d[sh1] = d_i1; } if (l_start < l_end) { float* od0 = sc.mla_att_o + size_t(sh0) * KV_LORA + lane * 4; float* od1 = sc.mla_att_o + size_t(sh1) * KV_LORA + lane * 4; #pragma unroll for (int c = 0; c < 4; ++c) { *reinterpret_cast(od0 + c * 128) = make_float4( o_acc_0[c * 4 + 0], o_acc_0[c * 4 + 1], o_acc_0[c * 4 + 2], o_acc_0[c * 4 + 3] ); *reinterpret_cast(od1 + c * 128) = make_float4( o_acc_1[c * 4 + 0], o_acc_1[c * 4 + 1], o_acc_1[c * 4 + 2], o_acc_1[c * 4 + 3] ); } } } grid_barrier(local_epoch); // ---------------- Phase M5: 128-Block Parallel Split-K Combine + Un-absorption (c_out @ W_uv) ---------------- if (bid < 128) { const int h = bid >> 2; const int g = bid & 3; float my_m = (tid < NUM_L_SPLITS) ? sc.mla_att_m[tid * MLA_H + h] : -1.0e30f; float my_d = (tid < NUM_L_SPLITS) ? sc.mla_att_d[tid * MLA_H + h] : 0.0f; float w_max = my_m; #pragma unroll for (int off = 16; off > 0; off >>= 1) { w_max = fmaxf(w_max, __shfl_xor_sync(0xffffffff, w_max, off)); } if (lane == 0) smem_red[warp_id] = w_max; __syncthreads(); if (tid < 8) { float gm = smem_red[tid]; #pragma unroll for (int off = 4; off > 0; off >>= 1) { gm = fmaxf(gm, __shfl_xor_sync(0xff, gm, off)); } if (tid == 0) smem_red[0] = gm; } __syncthreads(); float m_max = smem_red[0]; float my_w = (my_d > 0.0f) ? expf(my_m - m_max) : 0.0f; float wd_sum = warp_reduce_sum(my_d * my_w); if (lane == 0) smem_red[warp_id] = wd_sum; __syncthreads(); if (tid < 32) { float gd = (tid < 8) ? smem_red[tid] : 0.0f; gd = warp_reduce_sum(gd); if (tid == 0) smem_red[0] = 1.0f / gd; } __syncthreads(); float inv_d = smem_red[0]; if (tid < NUM_L_SPLITS) { smem_red[16 + tid] = my_w * inv_d; } __syncthreads(); float4 o_sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f); const int g_off = g * 128 + lane * 4; for (int s = warp_id; s < NUM_L_SPLITS; s += 8) { float ws = smem_red[16 + s]; if (ws > 0.0f) { float4 v4 = *reinterpret_cast(sc.mla_att_o + size_t(s * MLA_H + h) * KV_LORA + g_off); o_sum.x += v4.x * ws; o_sum.y += v4.y * ws; o_sum.z += v4.z * ws; o_sum.w += v4.w * ws; } } *reinterpret_cast(&smem_x[warp_id * 128 + lane * 4]) = o_sum; __syncthreads(); if (tid < 128) { float sum = smem_x[0 * 128 + tid] + smem_x[1 * 128 + tid] + smem_x[2 * 128 + tid] + smem_x[3 * 128 + tid] + smem_x[4 * 128 + tid] + smem_x[5 * 128 + tid] + smem_x[6 * 128 + tid] + smem_x[7 * 128 + tid]; smem_x[g * 128 + tid] = sum; } __syncthreads(); float acc_uv[4] = {0.0f, 0.0f, 0.0f, 0.0f}; const int n0 = h * KVB_DIM + QK_NOPE; w4a16_group_tile_128( mb.kv_b.w_q, mb.kv_b.scales, mb.kv_b.zeros, smem_x, 8192, n0, g, warp_id, lane, acc_uv ); flush_tile_128(sc.kda_o + h * V_HEAD, acc_uv, smem_red, tid, warp_id, lane); } grid_barrier(local_epoch); // ---------------- Phase M6: MLA o_proj W4A16 GEMV ---------------- for (int i = tid; i < 4096; i += BLOCK_THREADS) { smem_x[i] = round_bf16(sc.kda_o[i]); } __syncthreads(); { constexpr int TOTAL_O_GROUPS = 18 * 32; const int o_start = (eff_bid * TOTAL_O_GROUPS) / NUM_SMS; const int o_end = ((eff_bid + 1) * TOTAL_O_GROUPS) / NUM_SMS; float acc_o[4] = {0.0f, 0.0f, 0.0f, 0.0f}; int cur_oc = -1; int prev_oc = -1; for (int u = o_start; u < o_end; ++u) { int c_tile = u >> 5; int g = u & 31; if (c_tile != cur_oc) { if (cur_oc >= 0) { save_tile_128(acc_o, smem_red_prev, warp_id, lane); prev_oc = cur_oc; } cur_oc = c_tile; } w4a16_group_tile_128( mb.o_proj.w_q, mb.o_proj.scales, mb.o_proj.zeros, smem_x, HIDDEN, c_tile * 128, g, warp_id, lane, acc_o ); } flush_dual_tile_128( (cur_oc >= 0) ? (sc.attn_out_fp32 + cur_oc * 128) : nullptr, (prev_oc >= 0) ? (sc.attn_out_fp32 + prev_oc * 128) : nullptr, acc_o, smem_red, smem_red_prev, tid, warp_id, lane ); } grid_barrier(local_epoch); // ---------------- Phases M7, M8, M9: MLA MoE ---------------- run_moe_phases(local_epoch, mb.moe_norm, mb.moe, sc, smem_x, smem_red, smem_top_idx, smem_top_w); // ---------------- Final Epilogue: Write Final Hidden Output ---------------- if (bid < 9) { int i = bid * BLOCK_THREADS + tid; float h_old = bf16_to_f32(sc.res_h[i]); float m_out = round_bf16(sc.moe_out_fp32[i]); params.hidden_out[i] = f32_to_bf16(h_old + m_out); } } extern "C" void launch_kimi_megakernel(const void* params_ptr, cudaStream_t stream) { const KernelParams* p = reinterpret_cast(params_ptr); kimi_decode_megakernel<<>>(*p); } """ _CPP_SRC = r""" #include #include #include extern "C" void launch_kimi_megakernel(const void* params_ptr, cudaStream_t stream); struct QLinearPtrs { const void* w_q; const void* scales; const void* zeros; }; struct MoEPtrs { const void* router_w; QLinearPtrs gate; QLinearPtrs up; QLinearPtrs down; QLinearPtrs s_gate; QLinearPtrs s_up; QLinearPtrs s_down; }; struct KDABlockPtrs { const void* attn_norm; const void* moe_norm; QLinearPtrs q_proj; QLinearPtrs k_proj; QLinearPtrs v_proj; QLinearPtrs g_proj; const void* beta_w; const void* conv_w; QLinearPtrs o_proj; MoEPtrs moe; void* S; void* cq; void* ck; void* cv; }; struct MLABlockPtrs { const void* attn_norm; const void* moe_norm; QLinearPtrs q_proj; QLinearPtrs kv_a; QLinearPtrs kv_b; QLinearPtrs o_proj; MoEPtrs moe; const void* c_kv_src; const void* k_rope_src; void* c_kv_dst; void* k_rope_dst; int pos; }; struct ScratchPtrs { void* norm_x; void* attn_out_fp32; void* moe_out_fp32; void* res_h; void* kda_qkvg_fp32; void* kda_o; void* kda_beta; void* mla_q_kv_fp32; void* mla_q_abs; void* mla_q_rope; void* mla_att_m; void* mla_att_d; void* mla_att_o; void* router_logits; void* moe_gate_fp32; void* moe_up_fp32; }; struct KernelParams { KDABlockPtrs kda[3]; MLABlockPtrs mla; ScratchPtrs scratch; const void* hidden_in; void* hidden_out; }; static KernelParams g_params; void init_static_params(const std::vector& ptrs) { int idx = 0; auto next_ptr = [&]() -> void* { return reinterpret_cast(ptrs[idx++]); }; auto read_ql = [&]() -> QLinearPtrs { QLinearPtrs q; q.w_q = next_ptr(); q.scales = next_ptr(); q.zeros = next_ptr(); return q; }; auto read_moe = [&]() -> MoEPtrs { MoEPtrs m; m.router_w = next_ptr(); m.gate = read_ql(); m.up = read_ql(); m.down = read_ql(); m.s_gate = read_ql(); m.s_up = read_ql(); m.s_down = read_ql(); return m; }; for (int b = 0; b < 3; ++b) { g_params.kda[b].attn_norm = next_ptr(); g_params.kda[b].moe_norm = next_ptr(); g_params.kda[b].q_proj = read_ql(); g_params.kda[b].k_proj = read_ql(); g_params.kda[b].v_proj = read_ql(); g_params.kda[b].g_proj = read_ql(); g_params.kda[b].beta_w = next_ptr(); g_params.kda[b].conv_w = next_ptr(); g_params.kda[b].o_proj = read_ql(); g_params.kda[b].moe = read_moe(); } g_params.mla.attn_norm = next_ptr(); g_params.mla.moe_norm = next_ptr(); g_params.mla.q_proj = read_ql(); g_params.mla.kv_a = read_ql(); g_params.mla.kv_b = read_ql(); g_params.mla.o_proj = read_ql(); g_params.mla.moe = read_moe(); g_params.scratch.norm_x = next_ptr(); g_params.scratch.attn_out_fp32 = next_ptr(); g_params.scratch.moe_out_fp32 = next_ptr(); g_params.scratch.res_h = next_ptr(); g_params.scratch.kda_qkvg_fp32 = next_ptr(); g_params.scratch.kda_o = next_ptr(); g_params.scratch.kda_beta = next_ptr(); g_params.scratch.mla_q_kv_fp32 = next_ptr(); g_params.scratch.mla_q_abs = next_ptr(); g_params.scratch.mla_q_rope = next_ptr(); g_params.scratch.mla_att_m = next_ptr(); g_params.scratch.mla_att_d = next_ptr(); g_params.scratch.mla_att_o = next_ptr(); g_params.scratch.router_logits = next_ptr(); g_params.scratch.moe_gate_fp32 = next_ptr(); g_params.scratch.moe_up_fp32 = next_ptr(); } void run_step( torch::Tensor hidden_in, torch::Tensor hidden_out, torch::Tensor S0, torch::Tensor cq0, torch::Tensor ck0, torch::Tensor cv0, torch::Tensor S1, torch::Tensor cq1, torch::Tensor ck1, torch::Tensor cv1, torch::Tensor S2, torch::Tensor cq2, torch::Tensor ck2, torch::Tensor cv2, torch::Tensor c_kv_src, torch::Tensor k_rope_src, torch::Tensor c_kv_dst, torch::Tensor k_rope_dst, int64_t pos ) { g_params.hidden_in = hidden_in.data_ptr(); g_params.hidden_out = hidden_out.data_ptr(); g_params.kda[0].S = S0.data_ptr(); g_params.kda[0].cq = cq0.data_ptr(); g_params.kda[0].ck = ck0.data_ptr(); g_params.kda[0].cv = cv0.data_ptr(); g_params.kda[1].S = S1.data_ptr(); g_params.kda[1].cq = cq1.data_ptr(); g_params.kda[1].ck = ck1.data_ptr(); g_params.kda[1].cv = cv1.data_ptr(); g_params.kda[2].S = S2.data_ptr(); g_params.kda[2].cq = cq2.data_ptr(); g_params.kda[2].ck = ck2.data_ptr(); g_params.kda[2].cv = cv2.data_ptr(); g_params.mla.c_kv_src = c_kv_src.data_ptr(); g_params.mla.k_rope_src = k_rope_src.data_ptr(); g_params.mla.c_kv_dst = c_kv_dst.data_ptr(); g_params.mla.k_rope_dst = k_rope_dst.data_ptr(); g_params.mla.pos = static_cast(pos); cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); launch_kimi_megakernel(&g_params, stream); } """ _EXT = load_inline( name="kimi_megakernel_v5", cpp_sources=_CPP_SRC, cuda_sources=_CUDA_SRC, functions=["init_static_params", "run_step"], no_implicit_headers=True, extra_cuda_cflags=["-O3", "--use_fast_math"], ) class Model(nn.Module): def __init__(self, cfg: Config): super().__init__() self.cfg = cfg self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern) self._prepared_dev = None def reset_parameters(self): g = torch.Generator(device="cpu").manual_seed(1234) for mod in self.modules(): if isinstance(mod, (QuantLinear, QuantExperts)): mod.init_random(g) elif isinstance(mod, nn.Linear): nn.init.normal_(mod.weight, 0.0, 0.02, generator=g) elif isinstance(mod, KDA): nn.init.normal_(mod.conv_w, 0.0, 0.1, generator=g) self._prepared_dev = None def _prepare_buffers(self, dev: torch.device): self._norm_x = torch.empty(2304, dtype=torch.float32, device=dev) self._attn_out_fp32 = torch.empty(2304, dtype=torch.float32, device=dev) self._moe_out_fp32 = torch.empty(2304, dtype=torch.float32, device=dev) self._res_h = torch.empty(2 * 2304, dtype=torch.bfloat16, device=dev) self._kda_qkvg_fp32 = torch.empty(4 * 4096, dtype=torch.float32, device=dev) self._kda_o = torch.empty(4096, dtype=torch.float32, device=dev) self._kda_beta = torch.empty(32, dtype=torch.float32, device=dev) self._mla_q_kv_fp32 = torch.empty(53 * 128, dtype=torch.float32, device=dev) self._mla_q_abs = torch.empty(32 * 512, dtype=torch.float32, device=dev) self._mla_q_rope = torch.empty(32 * 64, dtype=torch.float32, device=dev) self._mla_att_m = torch.empty(188 * 32, dtype=torch.float32, device=dev) self._mla_att_d = torch.empty(188 * 32, dtype=torch.float32, device=dev) self._mla_att_o = torch.empty(188 * 32 * 512, dtype=torch.float32, device=dev) self._router_logits = torch.empty(64, dtype=torch.float32, device=dev) self._moe_gate_fp32 = torch.empty(9 * 1024, dtype=torch.float32, device=dev) self._moe_up_fp32 = torch.empty(9 * 1024, dtype=torch.float32, device=dev) self._c_kv_cap = torch.empty((20480, 512), dtype=torch.bfloat16, device=dev) self._k_rope_cap = torch.empty((20480, 64), dtype=torch.bfloat16, device=dev) self._hidden_out = torch.empty(2304, dtype=torch.bfloat16, device=dev) ptrs = [] def add_ql(ql): ptrs.extend([ql.w_q.data_ptr(), ql.scales.data_ptr(), ql.zeros.data_ptr()]) def add_moe(moe): ptrs.append(moe.router.weight.data_ptr()) add_ql(moe.gate) add_ql(moe.up) add_ql(moe.down) add_ql(moe.s_gate) add_ql(moe.s_up) add_ql(moe.s_down) for b in range(3): blk = self.blocks[b] ptrs.append(blk.attn_norm.data_ptr()) ptrs.append(blk.moe_norm.data_ptr()) add_ql(blk.attn.q_proj) add_ql(blk.attn.k_proj) add_ql(blk.attn.v_proj) add_ql(blk.attn.g_proj) ptrs.append(blk.attn.beta_proj.weight.data_ptr()) ptrs.append(blk.attn.conv_w.data_ptr()) add_ql(blk.attn.o_proj) add_moe(blk.moe) blk3 = self.blocks[3] ptrs.append(blk3.attn_norm.data_ptr()) ptrs.append(blk3.moe_norm.data_ptr()) add_ql(blk3.attn.q_proj) add_ql(blk3.attn.kv_a) add_ql(blk3.attn.kv_b) add_ql(blk3.attn.o_proj) add_moe(blk3.moe) ptrs.extend([ self._norm_x.data_ptr(), self._attn_out_fp32.data_ptr(), self._moe_out_fp32.data_ptr(), self._res_h.data_ptr(), self._kda_qkvg_fp32.data_ptr(), self._kda_o.data_ptr(), self._kda_beta.data_ptr(), self._mla_q_kv_fp32.data_ptr(), self._mla_q_abs.data_ptr(), self._mla_q_rope.data_ptr(), self._mla_att_m.data_ptr(), self._mla_att_d.data_ptr(), self._mla_att_o.data_ptr(), self._router_logits.data_ptr(), self._moe_gate_fp32.data_ptr(), self._moe_up_fp32.data_ptr(), ]) _EXT.init_static_params(ptrs) self._prepared_dev = dev def _apply(self, fn): ret = super()._apply(fn) dev = self.blocks[0].attn_norm.device if dev.type == "cuda": self._prepare_buffers(dev) return ret def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False): ret = super().load_state_dict(state_dict, strict=strict, assign=assign) dev = self.blocks[0].attn_norm.device if dev.type == "cuda": self._prepare_buffers(dev) return ret def step(self, hidden: torch.Tensor, state: list): if self._prepared_dev != hidden.device: self._prepare_buffers(hidden.device) s0, s1, s2, s3 = state[0], state[1], state[2], state[3] c_kv_src = s3["c_kv"] k_rope_src = s3["k_rope"] pos = c_kv_src.shape[0] _EXT.run_step( hidden, self._hidden_out, s0["S"], s0["cq"], s0["ck"], s0["cv"], s1["S"], s1["cq"], s1["ck"], s1["cv"], s2["S"], s2["cq"], s2["ck"], s2["cv"], c_kv_src, k_rope_src, self._c_kv_cap, self._k_rope_cap, pos, ) s3["c_kv"] = self._c_kv_cap[: pos + 1] s3["k_rope"] = self._k_rope_cap[: pos + 1] return self._hidden_out, state def init_state(cfg: Config, context_len: int, seed: int = 0, device: str = "cuda"): g = torch.Generator(device=device).manual_seed(seed) state = [] for k in cfg.pattern: if k == "K": H, Dk, W = cfg.kda_heads, cfg.kda_head_dim, cfg.short_conv S = torch.randn(H, Dk, Dk, generator=g, device=device, dtype=torch.float32) * 0.05 cq = torch.randn(W - 1, H * Dk, generator=g, device=device, dtype=cfg.dtype) * 0.1 ck = torch.randn(W - 1, H * Dk, generator=g, device=device, dtype=cfg.dtype) * 0.1 cv = torch.randn(W - 1, H * Dk, generator=g, device=device, dtype=cfg.dtype) * 0.1 state.append({"S": S, "cq": cq, "ck": ck, "cv": cv}) else: c_kv = torch.randn(context_len, cfg.kv_lora, generator=g, device=device, dtype=cfg.dtype) * 0.1 k_rope = torch.randn(context_len, cfg.qk_rope, generator=g, device=device, dtype=cfg.dtype) * 0.1 state.append({"c_kv": c_kv, "k_rope": k_rope}) return state def init_token(cfg: Config, seed: int = 0, device: str = "cuda") -> torch.Tensor: g = torch.Generator(device=device).manual_seed(seed) return torch.randn(cfg.hidden, generator=g, device=device, dtype=cfg.dtype) * 0.2