KernelBench cuda · RTX PRO 6000
MegaQwen Decode Gemini 4 Argon (High)
5.40%geomean peak fraction across shapes
Gemini 4 Argon wrote a single cooperative CUDA megakernel (188 CTAs x 256 threads, grid.sync between phases) that runs every decode step for all four layers: RMSNorm and QKV GEMV, q/k-norm, RoPE and the KV-cache append, split-K flash decoding over 47 units per KV head, o_proj, SwiGLU and down_proj. It pins the first 80 MB of a packed weight buffer in L2 and caches the untimed prefill.
harnessagyagent session5h 48mtotal wall5h 59mcheck59sbenchmark10moutput tokens—gpu-lock wait4mgpu-lock held12mregimethroughput
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
No per-shape benchmark data archived for this run.
Kernel source (redacted)
from __future__ import annotations
import math
import os
import sys
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_venv_bin = os.path.dirname(sys.executable)
if _venv_bin and _venv_bin not in os.environ.get("PATH", ""):
os.environ["PATH"] = _venv_bin + os.pathsep + os.environ.get("PATH", "")
OP_TYPE = "megaqwen_decode"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000"]
HIDDEN = 1024
INTERMEDIATE = 3072
NUM_Q = 16
NUM_KV = 8
HEAD_DIM = 128
NUM_LAYERS = 4
EPS = 1e-6
CUDA_SRC = r"""
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cstdint>
#include <algorithm>
namespace cg = cooperative_groups;
constexpr int WARP_SIZE = 32;
constexpr int HIDDEN_SIZE = 1024;
constexpr int INTERMEDIATE_SIZE = 3072;
constexpr int NUM_Q_HEADS = 16;
constexpr int NUM_KV_HEADS = 8;
constexpr int HEAD_DIM = 128;
constexpr int Q_SIZE = NUM_Q_HEADS * HEAD_DIM; // 2048
constexpr int KV_SIZE = NUM_KV_HEADS * HEAD_DIM; // 1024
constexpr int QKV_SIZE = Q_SIZE + 2 * KV_SIZE; // 4096
constexpr int MAX_LAYERS = 8;
constexpr int DEC_NUM_BLOCKS = 188;
constexpr int DEC_BLOCK_SIZE = 256;
constexpr int DEC_NUM_WARPS = DEC_BLOCK_SIZE / WARP_SIZE; // 8
constexpr int MAX_SPLITS_PER_KV = 24; // ceil(188 / 8) = 24
constexpr float RMS_EPS = 1e-6f;
constexpr float LOG2E = 1.44269504088896340736f;
struct LayerPointers {
const __nv_bfloat16* input_ln;
const __nv_bfloat16* qkv_proj; // contiguous [4096, 1024]
const __nv_bfloat16* q_norm;
const __nv_bfloat16* k_norm;
const __nv_bfloat16* o_proj;
const __nv_bfloat16* post_ln;
const __nv_bfloat16* gate_proj;
const __nv_bfloat16* up_proj;
const __nv_bfloat16* down_proj;
__nv_bfloat16* k_cache;
__nv_bfloat16* v_cache;
};
struct DecodeParams {
LayerPointers layers[MAX_LAYERS];
};
struct uint8_vec {
uint4 a;
uint4 b;
};
// 256-bit (32B = 16 bf16) load with L1::no_allocate and L2::evict_last (LDG.E.NA.ELL2.256.CONSTANT)
__device__ __forceinline__ uint8_vec ld_weight_u8(const void* ptr) {
uint8_vec ret;
asm volatile(
"ld.global.nc.L1::no_allocate.L2::evict_last.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
: "=r"(ret.a.x), "=r"(ret.a.y), "=r"(ret.a.z), "=r"(ret.a.w),
"=r"(ret.b.x), "=r"(ret.b.y), "=r"(ret.b.z), "=r"(ret.b.w)
: "l"(ptr)
);
return ret;
}
// 128-bit (16B = 8 bf16) load with L1::no_allocate (LDG.E.NA.128.CONSTANT)
__device__ __forceinline__ uint4 ld_weight_u4(const void* ptr) {
uint4 ret;
asm volatile(
"ld.global.nc.L1::no_allocate.v4.b32 {%0, %1, %2, %3}, [%4];"
: "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w)
: "l"(ptr)
);
return ret;
}
// 128-bit (16B = 8 bf16) KV cache streaming load with L2::evict_first (LDG.E.EF.128)
__device__ __forceinline__ uint4 ld_kv_stream_u4(const void* ptr) {
uint4 ret;
asm volatile(
"ld.global.cs.v4.b32 {%0, %1, %2, %3}, [%4];"
: "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w)
: "l"(ptr)
);
return ret;
}
__device__ __forceinline__ float warp_reduce_sum(float val) {
#pragma unroll
for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
val += __shfl_down_sync(0xffffffffu, val, offset);
}
return val;
}
__device__ __forceinline__ float silu_fn(float x) {
return x / (1.0f + __expf(-x));
}
__device__ __forceinline__ float dot8_bf16_f32(uint4 w, const float a[8]) {
const __nv_bfloat162* wb = reinterpret_cast<const __nv_bfloat162*>(&w);
float2 f0 = __bfloat1622float2(wb[0]);
float2 f1 = __bfloat1622float2(wb[1]);
float2 f2 = __bfloat1622float2(wb[2]);
float2 f3 = __bfloat1622float2(wb[3]);
return (f0.x * a[0] + f0.y * a[1]) +
(f1.x * a[2] + f1.y * a[3]) +
(f2.x * a[4] + f2.y * a[5]) +
(f3.x * a[6] + f3.y * a[7]);
}
__device__ __forceinline__ float dot16_bf16_f32(uint8_vec w, const float a[16]) {
return dot8_bf16_f32(w.a, a) + dot8_bf16_f32(w.b, a + 8);
}
__device__ __forceinline__ void unpack_bf16x8(uint4 u, float out[8]) {
const __nv_bfloat162* p = reinterpret_cast<const __nv_bfloat162*>(&u);
float2 f0 = __bfloat1622float2(p[0]);
float2 f1 = __bfloat1622float2(p[1]);
float2 f2 = __bfloat1622float2(p[2]);
float2 f3 = __bfloat1622float2(p[3]);
out[0] = f0.x; out[1] = f0.y;
out[2] = f1.x; out[3] = f1.y;
out[4] = f2.x; out[5] = f2.y;
out[6] = f3.x; out[7] = f3.y;
}
__device__ __forceinline__ void unpack_bf16x16(uint8_vec u, float out[16]) {
unpack_bf16x8(u.a, out);
unpack_bf16x8(u.b, out + 8);
}
__device__ __forceinline__ void attn_batch4_step8(
uint4 k0, uint4 k1, uint4 k2, uint4 k3,
uint4 v0, uint4 v1, uint4 v2, uint4 v3,
const float q0_reg[8], const float q1_reg[8],
float& max0, float& sum0, float acc0[8],
float& max1, float& sum1, float acc1[8]
) {
float s0_0 = dot8_bf16_f32(k0, q0_reg);
float s1_0 = dot8_bf16_f32(k0, q1_reg);
float s0_1 = dot8_bf16_f32(k1, q0_reg);
float s1_1 = dot8_bf16_f32(k1, q1_reg);
float s0_2 = dot8_bf16_f32(k2, q0_reg);
float s1_2 = dot8_bf16_f32(k2, q1_reg);
float s0_3 = dot8_bf16_f32(k3, q0_reg);
float s1_3 = dot8_bf16_f32(k3, q1_reg);
#pragma unroll
for (int offset = 8; offset > 0; offset >>= 1) {
s0_0 += __shfl_xor_sync(0xffffffffu, s0_0, offset);
s1_0 += __shfl_xor_sync(0xffffffffu, s1_0, offset);
s0_1 += __shfl_xor_sync(0xffffffffu, s0_1, offset);
s1_1 += __shfl_xor_sync(0xffffffffu, s1_1, offset);
s0_2 += __shfl_xor_sync(0xffffffffu, s0_2, offset);
s1_2 += __shfl_xor_sync(0xffffffffu, s1_2, offset);
s0_3 += __shfl_xor_sync(0xffffffffu, s0_3, offset);
s1_3 += __shfl_xor_sync(0xffffffffu, s1_3, offset);
}
// Compute online softmax weights and rescale accumulators BEFORE touching v0..v3
float old_max0 = max0;
max0 = fmaxf(max0, fmaxf(fmaxf(s0_0, s0_1), fmaxf(s0_2, s0_3)));
float exp_diff0 = exp2f(old_max0 - max0);
float w0_0 = exp2f(s0_0 - max0);
float w0_1 = exp2f(s0_1 - max0);
float w0_2 = exp2f(s0_2 - max0);
float w0_3 = exp2f(s0_3 - max0);
sum0 = sum0 * exp_diff0 + (w0_0 + w0_1 + w0_2 + w0_3);
float old_max1 = max1;
max1 = fmaxf(max1, fmaxf(fmaxf(s1_0, s1_1), fmaxf(s1_2, s1_3)));
float exp_diff1 = exp2f(old_max1 - max1);
float w1_0 = exp2f(s1_0 - max1);
float w1_1 = exp2f(s1_1 - max1);
float w1_2 = exp2f(s1_2 - max1);
float w1_3 = exp2f(s1_3 - max1);
sum1 = sum1 * exp_diff1 + (w1_0 + w1_1 + w1_2 + w1_3);
#pragma unroll
for (int j = 0; j < 8; j++) {
acc0[j] *= exp_diff0;
acc1[j] *= exp_diff1;
}
float vf[8];
unpack_bf16x8(v0, vf);
#pragma unroll
for (int j = 0; j < 8; j++) {
acc0[j] += w0_0 * vf[j];
acc1[j] += w1_0 * vf[j];
}
unpack_bf16x8(v1, vf);
#pragma unroll
for (int j = 0; j < 8; j++) {
acc0[j] += w0_1 * vf[j];
acc1[j] += w1_1 * vf[j];
}
unpack_bf16x8(v2, vf);
#pragma unroll
for (int j = 0; j < 8; j++) {
acc0[j] += w0_2 * vf[j];
acc1[j] += w1_2 * vf[j];
}
unpack_bf16x8(v3, vf);
#pragma unroll
for (int j = 0; j < 8; j++) {
acc0[j] += w0_3 * vf[j];
acc1[j] += w1_3 * vf[j];
}
}
__device__ __forceinline__ void attn_batch2_step8(
uint4 k0, uint4 k1,
uint4 v0, uint4 v1,
const float q0_reg[8], const float q1_reg[8],
float& max0, float& sum0, float acc0[8],
float& max1, float& sum1, float acc1[8]
) {
float s0_0 = dot8_bf16_f32(k0, q0_reg);
float s1_0 = dot8_bf16_f32(k0, q1_reg);
float s0_1 = dot8_bf16_f32(k1, q0_reg);
float s1_1 = dot8_bf16_f32(k1, q1_reg);
#pragma unroll
for (int offset = 8; offset > 0; offset >>= 1) {
s0_0 += __shfl_xor_sync(0xffffffffu, s0_0, offset);
s1_0 += __shfl_xor_sync(0xffffffffu, s1_0, offset);
s0_1 += __shfl_xor_sync(0xffffffffu, s0_1, offset);
s1_1 += __shfl_xor_sync(0xffffffffu, s1_1, offset);
}
float old_max0 = max0;
max0 = fmaxf(max0, fmaxf(s0_0, s0_1));
float exp_diff0 = exp2f(old_max0 - max0);
float w0_0 = exp2f(s0_0 - max0);
float w0_1 = exp2f(s0_1 - max0);
sum0 = sum0 * exp_diff0 + (w0_0 + w0_1);
float old_max1 = max1;
max1 = fmaxf(max1, fmaxf(s1_0, s1_1));
float exp_diff1 = exp2f(old_max1 - max1);
float w1_0 = exp2f(s1_0 - max1);
float w1_1 = exp2f(s1_1 - max1);
sum1 = sum1 * exp_diff1 + (w1_0 + w1_1);
float vf[8];
unpack_bf16x8(v0, vf);
#pragma unroll
for (int j = 0; j < 8; j++) {
acc0[j] = acc0[j] * exp_diff0 + w0_0 * vf[j];
acc1[j] = acc1[j] * exp_diff1 + w1_0 * vf[j];
}
unpack_bf16x8(v1, vf);
#pragma unroll
for (int j = 0; j < 8; j++) {
acc0[j] += w0_1 * vf[j];
acc1[j] += w1_1 * vf[j];
}
}
__device__ __forceinline__ void attn_single_step8(
uint4 k0, uint4 v0,
const float q0_reg[8], const float q1_reg[8],
float& max0, float& sum0, float acc0[8],
float& max1, float& sum1, float acc1[8],
uint32_t mask
) {
float s0 = dot8_bf16_f32(k0, q0_reg);
float s1 = dot8_bf16_f32(k0, q1_reg);
#pragma unroll
for (int offset = 8; offset > 0; offset >>= 1) {
s0 += __shfl_xor_sync(mask, s0, offset);
s1 += __shfl_xor_sync(mask, s1, offset);
}
float old_max0 = max0;
max0 = fmaxf(max0, s0);
float exp_diff0 = exp2f(old_max0 - max0);
float w0 = exp2f(s0 - max0);
sum0 = sum0 * exp_diff0 + w0;
float old_max1 = max1;
max1 = fmaxf(max1, s1);
float exp_diff1 = exp2f(old_max1 - max1);
float w1 = exp2f(s1 - max1);
sum1 = sum1 * exp_diff1 + w1;
float vf[8];
unpack_bf16x8(v0, vf);
#pragma unroll
for (int j = 0; j < 8; j++) {
acc0[j] = acc0[j] * exp_diff0 + w0 * vf[j];
acc1[j] = acc1[j] * exp_diff1 + w1 * vf[j];
}
}
__global__ void __launch_bounds__(DEC_BLOCK_SIZE, 1)
megaqwen_decode_kernel(
DecodeParams params,
const __nv_bfloat16* __restrict__ x_rand_stream, // [n_steps, 1024]
const __nv_bfloat16* __restrict__ hidden_in, // [1024]
__nv_bfloat16* __restrict__ hidden_out, // [1024]
float* __restrict__ g_residual, // [1024]
float* __restrict__ g_activations, // [1024]
float* __restrict__ g_qkv, // [4096]
float* __restrict__ g_attn_out, // [2048]
float* __restrict__ g_mlp_intermediate, // [3072]
float* __restrict__ g_split_max, // [16 * 24]
float* __restrict__ g_split_sum, // [16 * 24]
float* __restrict__ g_split_acc, // [16 * 24 * 128]
unsigned int* __restrict__ g_attn_lock, // [8]
int num_layers,
int start_pos,
int n_steps,
int max_seq_len,
float attn_scale
) {
cg::grid_group grid = cg::this_grid();
const int block_id = blockIdx.x;
const int tid = threadIdx.x;
const int warp_id = tid >> 5;
const int lane_id = tid & 31;
// Quarter-block coordinates (64 threads = 2 warps per quarter-block, 752 total across grid)
const int qgroup = tid >> 6; // 0..3
const int qtid = tid & 63; // 0..63
const int qwarp = (tid >> 5) & 1; // 0..1
const int global_qblock = (block_id << 2) + qgroup; // 0..751
// Half-block coordinates (128 threads = 4 warps per half-block, 376 total across grid)
const int hgroup = tid >> 7; // 0..1
const int htid = tid & 127; // 0..127
const int hwarp = (tid >> 5) & 3; // 0..3
const int global_hblock = (block_id << 1) + hgroup; // 0..375
// Half-warp coordinates for attention (16 threads per half-warp, 16 half-warps per block)
const int sub_id = tid >> 4; // 0..15
const int sub_lane = tid & 15; // 0..15
const int sub_warp = (tid >> 4) & 1; // 0..1 within warp
const uint32_t half_mask = sub_warp ? 0xffff0000u : 0x0000ffffu;
// Symmetrical 47-half-block attention coordinates (all 8 KV heads get 23 full blocks + 1 half-block = 24 splits)
const bool is_half_split = (block_id >= 184);
const int split_in_block = is_half_split ? hgroup : 0;
const int attn_kv_head = is_half_split ? (((block_id - 184) << 1) + hgroup) : (block_id & 7);
const int attn_split_idx = is_half_split ? 23 : (block_id >> 3);
const int sub_start = is_half_split ? (hgroup * 8) : 0;
const int num_subs = is_half_split ? 8 : 16;
const int local_sub = sub_id - sub_start;
const int local_warp = is_half_split ? hwarp : warp_id;
const int local_tid = is_half_split ? htid : tid;
const float q_scale_log2e = attn_scale * LOG2E;
// Shared memory
__shared__ __align__(16) float s_vec[HIDDEN_SIZE];
__shared__ float s_red[DEC_NUM_WARPS];
__shared__ float s_qpart_a[4][2][8]; // for quarter-block reductions (up to 6 rows)
__shared__ float s_qpart_b[4][2][8]; // second buffer for Gate+Up
__shared__ float s_hpart[2][4][4]; // for half-block reductions (up to 3 rows)
// Shared memory for attention (16 half-warps per block)
__shared__ __align__(16) float s_q[2][2][HEAD_DIM];
__shared__ __align__(16) float s_k_cur[HEAD_DIM];
__shared__ __align__(16) float s_v_cur[HEAD_DIM];
__shared__ float s_w_max[2][16];
__shared__ float s_w_sum[2][16];
__shared__ __align__(16) float s_w_acc[2][16][HEAD_DIM];
__shared__ int s_is_last_block[2];
// Precompute RoPE inverse frequencies in registers for this lane
const float inv_freq0 = 1.0f / powf(10000.0f, float(lane_id) / 64.0f);
const float inv_freq1 = 1.0f / powf(10000.0f, float(lane_id + 32) / 64.0f);
// Precompute row ranges for QKV (4096 rows across 752 quarter-blocks: 336 get 6, 416 get 5)
const int qkv_rows = (global_qblock < 336) ? 6 : 5;
const int qkv_start = (global_qblock < 336) ? (global_qblock * 6) : (336 * 6 + (global_qblock - 336) * 5);
// Precompute row ranges for O-proj and Down-proj (1024 rows across 376 half-blocks: 272 get 3, 104 get 2)
const int hid_rows = (global_hblock < 272) ? 3 : 2;
const int hid_start = (global_hblock < 272) ? (global_hblock * 3) : (272 * 3 + (global_hblock - 272) * 2);
// Precompute row ranges for Gate+Up (3072 rows across 752 quarter-blocks: 64 get 5, 688 get 4)
const int mlp_rows = (global_qblock < 64) ? 5 : 4;
const int mlp_start = (global_qblock < 64) ? (global_qblock * 5) : (64 * 5 + (global_qblock - 64) * 4);
for (int step = 0; step < n_steps; step++) {
const int position = start_pos + step;
const __nv_bfloat16* x_rand_step = x_rand_stream + (size_t)step * HIDDEN_SIZE;
float cos0, sin0, cos1, sin1;
sincosf(float(position) * inv_freq0, &sin0, &cos0);
sincosf(float(position) * inv_freq1, &sin1, &cos1);
for (int layer = 0; layer < num_layers; layer++) {
const LayerPointers& lw = params.layers[layer];
// =================================================================
// PHASE 1: Redundant Input RMSNorm + QKV Projection (256-bit .ELL2)
// =================================================================
{
float local_sq = 0.0f;
int idx4 = tid * 4;
if (layer == 0) {
const __nv_bfloat16* h_src = (step == 0) ? hidden_in : hidden_out;
uint2 xr_u2 = *reinterpret_cast<const uint2*>(x_rand_step + idx4);
uint2 hb_u2 = *reinterpret_cast<const uint2*>(h_src + idx4);
const __nv_bfloat16* xr_p = reinterpret_cast<const __nv_bfloat16*>(&xr_u2);
const __nv_bfloat16* hb_p = reinterpret_cast<const __nv_bfloat16*>(&hb_u2);
float4 v4;
#pragma unroll
for (int j = 0; j < 4; j++) {
__nv_bfloat16 mixed = __float2bfloat16(0.5f * __bfloat162float(xr_p[j]) + 0.5f * __bfloat162float(hb_p[j]));
float vf = __bfloat162float(mixed);
(&v4.x)[j] = vf;
local_sq += vf * vf;
}
*reinterpret_cast<float4*>(s_vec + idx4) = v4;
if (block_id == 0) {
*reinterpret_cast<float4*>(g_residual + idx4) = v4;
}
} else {
uint2 hb_u2 = *reinterpret_cast<const uint2*>(hidden_out + idx4);
const __nv_bfloat16* hb_p = reinterpret_cast<const __nv_bfloat16*>(&hb_u2);
float4 v4;
v4.x = __bfloat162float(hb_p[0]);
v4.y = __bfloat162float(hb_p[1]);
v4.z = __bfloat162float(hb_p[2]);
v4.w = __bfloat162float(hb_p[3]);
local_sq = v4.x * v4.x + v4.y * v4.y + v4.z * v4.z + v4.w * v4.w;
*reinterpret_cast<float4*>(s_vec + idx4) = v4;
if (block_id == 0) {
*reinterpret_cast<float4*>(g_residual + idx4) = v4;
}
}
const int k_base = qtid * 16;
uint8_vec ln_u8 = ld_weight_u8(lw.input_ln + k_base);
local_sq = warp_reduce_sum(local_sq);
if (lane_id == 0) s_red[warp_id] = local_sq;
__syncthreads();
if (warp_id == 0) {
float s = (lane_id < DEC_NUM_WARPS) ? s_red[lane_id] : 0.0f;
s = warp_reduce_sum(s);
if (lane_id == 0) {
s_red[0] = rsqrtf(s * (1.0f / float(HIDDEN_SIZE)) + RMS_EPS);
}
}
__syncthreads();
const __nv_bfloat16* qkv_base = lw.qkv_proj + qkv_start * HIDDEN_SIZE + k_base;
uint8_vec qw[6];
#pragma unroll
for (int r = 0; r < 6; r++) {
qw[r] = ld_weight_u8(qkv_base + r * HIDDEN_SIZE);
}
const float rstd = s_red[0];
float4 sv0 = *reinterpret_cast<const float4*>(s_vec + k_base);
float4 sv1 = *reinterpret_cast<const float4*>(s_vec + k_base + 4);
float4 sv2 = *reinterpret_cast<const float4*>(s_vec + k_base + 8);
float4 sv3 = *reinterpret_cast<const float4*>(s_vec + k_base + 12);
float lw_f[16];
unpack_bf16x16(ln_u8, lw_f);
float a[16];
a[0] = sv0.x * rstd * lw_f[0];
a[1] = sv0.y * rstd * lw_f[1];
a[2] = sv0.z * rstd * lw_f[2];
a[3] = sv0.w * rstd * lw_f[3];
a[4] = sv1.x * rstd * lw_f[4];
a[5] = sv1.y * rstd * lw_f[5];
a[6] = sv1.z * rstd * lw_f[6];
a[7] = sv1.w * rstd * lw_f[7];
a[8] = sv2.x * rstd * lw_f[8];
a[9] = sv2.y * rstd * lw_f[9];
a[10] = sv2.z * rstd * lw_f[10];
a[11] = sv2.w * rstd * lw_f[11];
a[12] = sv3.x * rstd * lw_f[12];
a[13] = sv3.y * rstd * lw_f[13];
a[14] = sv3.z * rstd * lw_f[14];
a[15] = sv3.w * rstd * lw_f[15];
#pragma unroll
for (int r = 0; r < 6; r++) {
float p = warp_reduce_sum(dot16_bf16_f32(qw[r], a));
if (lane_id == 0) {
s_qpart_a[qgroup][qwarp][r] = p;
}
}
__syncthreads();
if (qwarp == 0 && lane_id < qkv_rows) {
float sum = s_qpart_a[qgroup][0][lane_id] + s_qpart_a[qgroup][1][lane_id];
g_qkv[qkv_start + lane_id] = sum;
}
}
grid.sync();
// =================================================================
// PHASE 2: Fused Q/K RMSNorm + RoPE + KV Cache + 47-Unit Split Attn
// =================================================================
{
const int hist_len = position;
const int unit_chunk = max(8, ((hist_len + 46) / 47 + 7) & ~7);
const int full_chunk = unit_chunk << 1;
const int active_splits = min(24, max(1, (hist_len + full_chunk - 1) / full_chunk));
if (attn_split_idx < active_splits) {
const float* g_q = g_qkv;
const float* g_k = g_qkv + Q_SIZE;
const float* g_v = g_qkv + Q_SIZE + KV_SIZE;
// Local Warps 0 & 1: Q RMSNorm + RoPE + pre-scale by attn_scale * LOG2E
if (local_warp < 2) {
const int qh = attn_kv_head * 2 + local_warp;
float* sq_dst = s_q[split_in_block][local_warp];
const float* q_in = g_q + qh * HEAD_DIM;
float q0 = q_in[lane_id];
float q1 = q_in[lane_id + 32];
float q2 = q_in[lane_id + 64];
float q3 = q_in[lane_id + 96];
float wn0 = __bfloat162float(__ldg(lw.q_norm + lane_id));
float wn1 = __bfloat162float(__ldg(lw.q_norm + lane_id + 32));
float wn2 = __bfloat162float(__ldg(lw.q_norm + lane_id + 64));
float wn3 = __bfloat162float(__ldg(lw.q_norm + lane_id + 96));
float sq = q0 * q0 + q1 * q1 + q2 * q2 + q3 * q3;
sq = warp_reduce_sum(sq);
float scale = rsqrtf(sq * (1.0f / float(HEAD_DIM)) + RMS_EPS);
scale = __shfl_sync(0xffffffffu, scale, 0) * q_scale_log2e;
q0 *= scale * wn0;
q1 *= scale * wn1;
q2 *= scale * wn2;
q3 *= scale * wn3;
sq_dst[lane_id] = q0 * cos0 - q2 * sin0;
sq_dst[lane_id + 64] = q0 * sin0 + q2 * cos0;
sq_dst[lane_id + 32] = q1 * cos1 - q3 * sin1;
sq_dst[lane_id + 96] = q1 * sin1 + q3 * cos1;
} else if (attn_split_idx == 0 && warp_id == 2) {
const float* k_in = g_k + attn_kv_head * HEAD_DIM;
float k0 = k_in[lane_id];
float k1 = k_in[lane_id + 32];
float k2 = k_in[lane_id + 64];
float k3 = k_in[lane_id + 96];
float kn0 = __bfloat162float(__ldg(lw.k_norm + lane_id));
float kn1 = __bfloat162float(__ldg(lw.k_norm + lane_id + 32));
float kn2 = __bfloat162float(__ldg(lw.k_norm + lane_id + 64));
float kn3 = __bfloat162float(__ldg(lw.k_norm + lane_id + 96));
float sq = k0 * k0 + k1 * k1 + k2 * k2 + k3 * k3;
sq = warp_reduce_sum(sq);
float scale = rsqrtf(sq * (1.0f / float(HEAD_DIM)) + RMS_EPS);
scale = __shfl_sync(0xffffffffu, scale, 0);
k0 *= scale * kn0;
k1 *= scale * kn1;
k2 *= scale * kn2;
k3 *= scale * kn3;
float kf0 = k0 * cos0 - k2 * sin0;
float kf2 = k0 * sin0 + k2 * cos0;
float kf1 = k1 * cos1 - k3 * sin1;
float kf3 = k1 * sin1 + k3 * cos1;
__nv_bfloat16 kb0 = __float2bfloat16(kf0);
__nv_bfloat16 kb1 = __float2bfloat16(kf1);
__nv_bfloat16 kb2 = __float2bfloat16(kf2);
__nv_bfloat16 kb3 = __float2bfloat16(kf3);
size_t cache_offset = ((size_t)attn_kv_head * max_seq_len + position) * HEAD_DIM;
__nv_bfloat16* kc_ptr = lw.k_cache + cache_offset;
kc_ptr[lane_id] = kb0;
kc_ptr[lane_id + 32] = kb1;
kc_ptr[lane_id + 64] = kb2;
kc_ptr[lane_id + 96] = kb3;
s_k_cur[lane_id] = __bfloat162float(kb0);
s_k_cur[lane_id + 32] = __bfloat162float(kb1);
s_k_cur[lane_id + 64] = __bfloat162float(kb2);
s_k_cur[lane_id + 96] = __bfloat162float(kb3);
} else if (attn_split_idx == 0 && warp_id == 3) {
const int d4 = lane_id * 4;
float4 vf = *reinterpret_cast<const float4*>(g_v + attn_kv_head * HEAD_DIM + d4);
__nv_bfloat162 vb01 = __floats2bfloat162_rn(vf.x, vf.y);
__nv_bfloat162 vb23 = __floats2bfloat162_rn(vf.z, vf.w);
uint2 v_u2;
v_u2.x = *reinterpret_cast<const uint32_t*>(&vb01);
v_u2.y = *reinterpret_cast<const uint32_t*>(&vb23);
size_t cache_offset = ((size_t)attn_kv_head * max_seq_len + position) * HEAD_DIM + d4;
*reinterpret_cast<uint2*>(lw.v_cache + cache_offset) = v_u2;
float2 f01 = __bfloat1622float2(vb01);
float2 f23 = __bfloat1622float2(vb23);
*reinterpret_cast<float4*>(s_v_cur + d4) = make_float4(f01.x, f01.y, f23.x, f23.y);
}
__syncthreads();
const int d_base = sub_lane * 8;
float4 q0_a = *reinterpret_cast<const float4*>(s_q[split_in_block][0] + d_base);
float4 q0_b = *reinterpret_cast<const float4*>(s_q[split_in_block][0] + d_base + 4);
float4 q1_a = *reinterpret_cast<const float4*>(s_q[split_in_block][1] + d_base);
float4 q1_b = *reinterpret_cast<const float4*>(s_q[split_in_block][1] + d_base + 4);
float q0_reg[8] = {q0_a.x, q0_a.y, q0_a.z, q0_a.w, q0_b.x, q0_b.y, q0_b.z, q0_b.w};
float q1_reg[8] = {q1_a.x, q1_a.y, q1_a.z, q1_a.w, q1_b.x, q1_b.y, q1_b.z, q1_b.w};
float max0 = -INFINITY, sum0 = 0.0f;
float acc0[8] = {0.0f};
float max1 = -INFINITY, sum1 = 0.0f;
float acc1[8] = {0.0f};
if (attn_split_idx == 0 && sub_id == 0) {
float s0 = 0.0f, s1 = 0.0f;
#pragma unroll
for (int j = 0; j < 8; j++) {
float kv = s_k_cur[d_base + j];
s0 += q0_reg[j] * kv;
s1 += q1_reg[j] * kv;
float vv = s_v_cur[d_base + j];
acc0[j] = vv;
acc1[j] = vv;
}
#pragma unroll
for (int offset = 8; offset > 0; offset >>= 1) {
s0 += __shfl_xor_sync(0x0000ffffu, s0, offset);
s1 += __shfl_xor_sync(0x0000ffffu, s1, offset);
}
max0 = s0;
max1 = s1;
sum0 = 1.0f;
sum1 = 1.0f;
}
const int split_span = is_half_split ? unit_chunk : full_chunk;
const int t_start = min(hist_len, attn_split_idx * full_chunk);
const int t_end = min(hist_len, t_start + split_span);
const __nv_bfloat16* k_head = lw.k_cache + (size_t)attn_kv_head * max_seq_len * HEAD_DIM;
const __nv_bfloat16* v_head = lw.v_cache + (size_t)attn_kv_head * max_seq_len * HEAD_DIM;
const int step4 = num_subs << 2; // 64 for full block, 32 for half block
const int step2 = num_subs << 1; // 32 for full block, 16 for half block
const int stride1 = num_subs * HEAD_DIM;
const int stride2 = (num_subs * 2) * HEAD_DIM;
const int stride3 = (num_subs * 3) * HEAD_DIM;
const int stride4 = step4 * HEAD_DIM;
int t_base = t_start;
const __nv_bfloat16* k_ptr = k_head + (size_t)(t_start + local_sub) * HEAD_DIM + d_base;
const __nv_bfloat16* v_ptr = v_head + (size_t)(t_start + local_sub) * HEAD_DIM + d_base;
// Uniform 4x-batched loop (512B contiguous per warp instruction, zero warp divergence)
for (; t_base + step4 <= t_end; t_base += step4) {
uint4 k0 = ld_kv_stream_u4(k_ptr);
uint4 k1 = ld_kv_stream_u4(k_ptr + stride1);
uint4 k2 = ld_kv_stream_u4(k_ptr + stride2);
uint4 k3 = ld_kv_stream_u4(k_ptr + stride3);
uint4 v0 = ld_kv_stream_u4(v_ptr);
uint4 v1 = ld_kv_stream_u4(v_ptr + stride1);
uint4 v2 = ld_kv_stream_u4(v_ptr + stride2);
uint4 v3 = ld_kv_stream_u4(v_ptr + stride3);
attn_batch4_step8(k0, k1, k2, k3, v0, v1, v2, v3, q0_reg, q1_reg, max0, sum0, acc0, max1, sum1, acc1);
k_ptr += stride4;
v_ptr += stride4;
}
// Uniform 2x-batched step
if (t_base + step2 <= t_end) {
uint4 k0 = ld_kv_stream_u4(k_ptr);
uint4 k1 = ld_kv_stream_u4(k_ptr + stride1);
uint4 v0 = ld_kv_stream_u4(v_ptr);
uint4 v1 = ld_kv_stream_u4(v_ptr + stride1);
attn_batch2_step8(k0, k1, v0, v1, q0_reg, q1_reg, max0, sum0, acc0, max1, sum1, acc1);
t_base += step2;
}
// Remaining < step2 positions (at most 1-2 iterations on the last split only)
for (int t = t_base + local_sub; t < t_end; t += num_subs) {
uint4 k0 = ld_kv_stream_u4(k_head + (size_t)t * HEAD_DIM + d_base);
uint4 v0 = ld_kv_stream_u4(v_head + (size_t)t * HEAD_DIM + d_base);
attn_single_step8(k0, v0, q0_reg, q1_reg, max0, sum0, acc0, max1, sum1, acc1, half_mask);
}
if (sub_lane == 0) {
s_w_max[0][sub_id] = max0;
s_w_sum[0][sub_id] = sum0;
s_w_max[1][sub_id] = max1;
s_w_sum[1][sub_id] = sum1;
}
*reinterpret_cast<float4*>(&s_w_acc[0][sub_id][d_base]) = make_float4(acc0[0], acc0[1], acc0[2], acc0[3]);
*reinterpret_cast<float4*>(&s_w_acc[0][sub_id][d_base + 4]) = make_float4(acc0[4], acc0[5], acc0[6], acc0[7]);
*reinterpret_cast<float4*>(&s_w_acc[1][sub_id][d_base]) = make_float4(acc1[0], acc1[1], acc1[2], acc1[3]);
*reinterpret_cast<float4*>(&s_w_acc[1][sub_id][d_base + 4]) = make_float4(acc1[4], acc1[5], acc1[6], acc1[7]);
__syncthreads();
// Block / half-block reduction across num_subs half-warps
if (local_warp < 2) {
const int h_sub = local_warp; // 0 or 1
const int qh = attn_kv_head * 2 + h_sub;
const int d_base4 = lane_id * 4;
float wm_lane = (lane_id < num_subs) ? s_w_max[h_sub][sub_start + lane_id] : -INFINITY;
float ws_lane = (lane_id < num_subs) ? s_w_sum[h_sub][sub_start + lane_id] : 0.0f;
float b_max = wm_lane;
#pragma unroll
for (int offset = 8; offset > 0; offset >>= 1) {
b_max = fmaxf(b_max, __shfl_xor_sync(0xffffffffu, b_max, offset));
}
b_max = __shfl_sync(0xffffffffu, b_max, 0);
float sc_lane = (wm_lane > -INFINITY) ? exp2f(wm_lane - b_max) : 0.0f;
float b_sum = ws_lane * sc_lane;
#pragma unroll
for (int offset = 8; offset > 0; offset >>= 1) {
b_sum += __shfl_xor_sync(0xffffffffu, b_sum, offset);
}
b_sum = __shfl_sync(0xffffffffu, b_sum, 0);
float b_out[4] = {0.0f, 0.0f, 0.0f, 0.0f};
if (b_max > -INFINITY) {
for (int w = 0; w < num_subs; w++) {
float sc = __shfl_sync(0xffffffffu, sc_lane, w);
if (sc > 0.0f) {
float4 wa = *reinterpret_cast<const float4*>(&s_w_acc[h_sub][sub_start + w][d_base4]);
b_out[0] += wa.x * sc;
b_out[1] += wa.y * sc;
b_out[2] += wa.z * sc;
b_out[3] += wa.w * sc;
}
}
}
if (active_splits == 1) {
float inv_sum = 1.0f / b_sum;
*reinterpret_cast<float4*>(g_attn_out + qh * HEAD_DIM + d_base4) =
make_float4(b_out[0] * inv_sum, b_out[1] * inv_sum,
b_out[2] * inv_sum, b_out[3] * inv_sum);
} else {
const int slot = qh * MAX_SPLITS_PER_KV + attn_split_idx;
if (lane_id == 0) {
g_split_max[slot] = b_max;
g_split_sum[slot] = b_sum;
}
*reinterpret_cast<float4*>(g_split_acc + slot * HEAD_DIM + d_base4) =
make_float4(b_out[0], b_out[1], b_out[2], b_out[3]);
}
}
if (active_splits > 1) {
__syncthreads();
if (local_tid == 0) {
asm volatile("fence.acq_rel.gpu;" ::: "memory");
unsigned int arrived = atomicAdd(&g_attn_lock[attn_kv_head], 1u);
s_is_last_block[split_in_block] = (arrived == (unsigned int)(active_splits - 1)) ? 1 : 0;
}
__syncthreads();
const int h_sub = is_half_split ? (local_warp >> 1) : (local_warp >> 2);
const int w_sub = is_half_split ? (local_warp & 1) : (local_warp & 3);
const int num_red_warps = is_half_split ? 2 : 4;
const int qh = attn_kv_head * 2 + h_sub;
const int base_slot = qh * MAX_SPLITS_PER_KV;
const int d_base4 = lane_id * 4;
float g_sum = 1.0f;
if (s_is_last_block[split_in_block]) {
asm volatile("fence.acq_rel.gpu;" ::: "memory");
float sm_lane = (lane_id < active_splits) ? g_split_max[base_slot + lane_id] : -INFINITY;
float ss_lane = (lane_id < active_splits) ? g_split_sum[base_slot + lane_id] : 0.0f;
float g_max = sm_lane;
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
g_max = fmaxf(g_max, __shfl_xor_sync(0xffffffffu, g_max, offset));
}
float sc_lane = (sm_lane > -INFINITY) ? exp2f(sm_lane - g_max) : 0.0f;
g_sum = ss_lane * sc_lane;
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
g_sum += __shfl_xor_sync(0xffffffffu, g_sum, offset);
}
float part_acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
for (int s = w_sub; s < active_splits; s += num_red_warps) {
float sc = __shfl_sync(0xffffffffu, sc_lane, s);
if (sc > 0.0f) {
float4 sa = *reinterpret_cast<const float4*>(g_split_acc + (base_slot + s) * HEAD_DIM + d_base4);
part_acc[0] += sa.x * sc;
part_acc[1] += sa.y * sc;
part_acc[2] += sa.z * sc;
part_acc[3] += sa.w * sc;
}
}
*reinterpret_cast<float4*>(&s_w_acc[h_sub][sub_start + w_sub][d_base4]) =
make_float4(part_acc[0], part_acc[1], part_acc[2], part_acc[3]);
}
__syncthreads();
if (s_is_last_block[split_in_block]) {
if (w_sub == 0) {
float4 p0 = *reinterpret_cast<const float4*>(&s_w_acc[h_sub][sub_start + 0][d_base4]);
float4 p1 = *reinterpret_cast<const float4*>(&s_w_acc[h_sub][sub_start + 1][d_base4]);
float4 p2 = is_half_split ? make_float4(0.f, 0.f, 0.f, 0.f) : *reinterpret_cast<const float4*>(&s_w_acc[h_sub][sub_start + 2][d_base4]);
float4 p3 = is_half_split ? make_float4(0.f, 0.f, 0.f, 0.f) : *reinterpret_cast<const float4*>(&s_w_acc[h_sub][sub_start + 3][d_base4]);
float inv_sum = 1.0f / g_sum;
*reinterpret_cast<float4*>(g_attn_out + qh * HEAD_DIM + d_base4) =
make_float4((p0.x + p1.x + p2.x + p3.x) * inv_sum,
(p0.y + p1.y + p2.y + p3.y) * inv_sum,
(p0.z + p1.z + p2.z + p3.z) * inv_sum,
(p0.w + p1.w + p2.w + p3.w) * inv_sum);
}
if (local_tid == 0) {
g_attn_lock[attn_kv_head] = 0u;
}
}
}
}
}
grid.sync();
// =================================================================
// PHASE 3: O Projection (1024 x 2048, 256-bit .ELL2) + Residual Add
// =================================================================
{
const int k_base = htid * 16; // 128 * 16 = 2048 = Q_SIZE
const __nv_bfloat16* o_base = lw.o_proj + hid_start * Q_SIZE + k_base;
uint8_vec ow[3];
#pragma unroll
for (int r = 0; r < 3; r++) {
ow[r] = ld_weight_u8(o_base + r * Q_SIZE);
}
float4 a0 = *reinterpret_cast<const float4*>(g_attn_out + k_base);
float4 a1 = *reinterpret_cast<const float4*>(g_attn_out + k_base + 4);
float4 a2 = *reinterpret_cast<const float4*>(g_attn_out + k_base + 8);
float4 a3 = *reinterpret_cast<const float4*>(g_attn_out + k_base + 12);
float a[16] = {
a0.x, a0.y, a0.z, a0.w,
a1.x, a1.y, a1.z, a1.w,
a2.x, a2.y, a2.z, a2.w,
a3.x, a3.y, a3.z, a3.w
};
#pragma unroll
for (int r = 0; r < 3; r++) {
float p = warp_reduce_sum(dot16_bf16_f32(ow[r], a));
if (lane_id == 0) {
s_hpart[hgroup][hwarp][r] = p;
}
}
__syncthreads();
if (hwarp == 0 && lane_id < hid_rows) {
float sum = s_hpart[hgroup][0][lane_id] +
s_hpart[hgroup][1][lane_id] +
s_hpart[hgroup][2][lane_id] +
s_hpart[hgroup][3][lane_id];
int m = hid_start + lane_id;
g_activations[m] = sum + g_residual[m];
}
}
grid.sync();
// =================================================================
// PHASE 4: Redundant Post-Attn RMSNorm + Gate & Up Proj + SiLU
// =================================================================
{
int idx4 = tid * 4;
float4 v4 = *reinterpret_cast<const float4*>(g_activations + idx4);
*reinterpret_cast<float4*>(s_vec + idx4) = v4;
float local_sq = v4.x * v4.x + v4.y * v4.y + v4.z * v4.z + v4.w * v4.w;
const int k_base = qtid * 16;
uint8_vec ln_u8 = ld_weight_u8(lw.post_ln + k_base);
local_sq = warp_reduce_sum(local_sq);
if (lane_id == 0) s_red[warp_id] = local_sq;
__syncthreads();
if (warp_id == 0) {
float s = (lane_id < DEC_NUM_WARPS) ? s_red[lane_id] : 0.0f;
s = warp_reduce_sum(s);
if (lane_id == 0) {
s_red[0] = rsqrtf(s * (1.0f / float(HIDDEN_SIZE)) + RMS_EPS);
}
}
__syncthreads();
const __nv_bfloat16* gate_base = lw.gate_proj + mlp_start * HIDDEN_SIZE + k_base;
const __nv_bfloat16* up_base = lw.up_proj + mlp_start * HIDDEN_SIZE + k_base;
uint8_vec gw[5], uw[5];
#pragma unroll
for (int r = 0; r < 5; r++) {
gw[r] = ld_weight_u8(gate_base + r * HIDDEN_SIZE);
uw[r] = ld_weight_u8(up_base + r * HIDDEN_SIZE);
}
const float rstd = s_red[0];
float4 sv0 = *reinterpret_cast<const float4*>(s_vec + k_base);
float4 sv1 = *reinterpret_cast<const float4*>(s_vec + k_base + 4);
float4 sv2 = *reinterpret_cast<const float4*>(s_vec + k_base + 8);
float4 sv3 = *reinterpret_cast<const float4*>(s_vec + k_base + 12);
float lw_f[16];
unpack_bf16x16(ln_u8, lw_f);
float a[16];
a[0] = sv0.x * rstd * lw_f[0];
a[1] = sv0.y * rstd * lw_f[1];
a[2] = sv0.z * rstd * lw_f[2];
a[3] = sv0.w * rstd * lw_f[3];
a[4] = sv1.x * rstd * lw_f[4];
a[5] = sv1.y * rstd * lw_f[5];
a[6] = sv1.z * rstd * lw_f[6];
a[7] = sv1.w * rstd * lw_f[7];
a[8] = sv2.x * rstd * lw_f[8];
a[9] = sv2.y * rstd * lw_f[9];
a[10] = sv2.z * rstd * lw_f[10];
a[11] = sv2.w * rstd * lw_f[11];
a[12] = sv3.x * rstd * lw_f[12];
a[13] = sv3.y * rstd * lw_f[13];
a[14] = sv3.z * rstd * lw_f[14];
a[15] = sv3.w * rstd * lw_f[15];
#pragma unroll
for (int r = 0; r < 5; r++) {
float gp = warp_reduce_sum(dot16_bf16_f32(gw[r], a));
float up = warp_reduce_sum(dot16_bf16_f32(uw[r], a));
if (lane_id == 0) {
s_qpart_a[qgroup][qwarp][r] = gp;
s_qpart_b[qgroup][qwarp][r] = up;
}
}
__syncthreads();
if (qwarp == 0 && lane_id < mlp_rows) {
float g_sum = s_qpart_a[qgroup][0][lane_id] + s_qpart_a[qgroup][1][lane_id];
float u_sum = s_qpart_b[qgroup][0][lane_id] + s_qpart_b[qgroup][1][lane_id];
g_mlp_intermediate[mlp_start + lane_id] = silu_fn(g_sum) * u_sum;
}
}
grid.sync();
// =================================================================
// PHASE 5: Down Projection (1024 x 3072) + Residual Add
// =================================================================
{
const int k0_base = htid * 16; // 0..2047
const int k1_base = 2048 + htid * 8; // 2048..3071
const __nv_bfloat16* down_base = lw.down_proj + hid_start * INTERMEDIATE_SIZE;
uint8_vec dw0[3];
uint4 dw1[3];
#pragma unroll
for (int r = 0; r < 3; r++) {
const __nv_bfloat16* row_ptr = down_base + r * INTERMEDIATE_SIZE;
dw0[r] = ld_weight_u8(row_ptr + k0_base);
dw1[r] = ld_weight_u4(row_ptr + k1_base);
}
float4 m0 = *reinterpret_cast<const float4*>(g_mlp_intermediate + k0_base);
float4 m1 = *reinterpret_cast<const float4*>(g_mlp_intermediate + k0_base + 4);
float4 m2 = *reinterpret_cast<const float4*>(g_mlp_intermediate + k0_base + 8);
float4 m3 = *reinterpret_cast<const float4*>(g_mlp_intermediate + k0_base + 12);
float4 m4 = *reinterpret_cast<const float4*>(g_mlp_intermediate + k1_base);
float4 m5 = *reinterpret_cast<const float4*>(g_mlp_intermediate + k1_base + 4);
float a0[16] = {
m0.x, m0.y, m0.z, m0.w,
m1.x, m1.y, m1.z, m1.w,
m2.x, m2.y, m2.z, m2.w,
m3.x, m3.y, m3.z, m3.w
};
float a1[8] = {
m4.x, m4.y, m4.z, m4.w,
m5.x, m5.y, m5.z, m5.w
};
#pragma unroll
for (int r = 0; r < 3; r++) {
float p = dot16_bf16_f32(dw0[r], a0) + dot8_bf16_f32(dw1[r], a1);
p = warp_reduce_sum(p);
if (lane_id == 0) {
s_hpart[hgroup][hwarp][r] = p;
}
}
__syncthreads();
if (hwarp == 0 && lane_id < hid_rows) {
float sum = s_hpart[hgroup][0][lane_id] +
s_hpart[hgroup][1][lane_id] +
s_hpart[hgroup][2][lane_id] +
s_hpart[hgroup][3][lane_id];
int m = hid_start + lane_id;
hidden_out[m] = __float2bfloat16(sum + g_activations[m]);
}
}
grid.sync();
}
}
}
void launch_megaqwen_decode(
torch::Tensor x_rand_stream,
torch::Tensor hidden_in,
torch::Tensor hidden_out,
torch::Tensor flat_weights,
std::vector<torch::Tensor> k_caches,
std::vector<torch::Tensor> v_caches,
torch::Tensor workspace_f32,
torch::Tensor attn_locks,
int num_layers,
int start_pos,
int n_steps,
int max_seq_len,
float attn_scale
) {
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const __nv_bfloat16* w_base = reinterpret_cast<const __nv_bfloat16*>(flat_weights.data_ptr());
const size_t w_bytes = (size_t)flat_weights.numel() * sizeof(__nv_bfloat16);
static const void* s_last_w_base = nullptr;
if (w_base != nullptr && (const void*)w_base != s_last_w_base) {
constexpr size_t MAX_PERSIST_BYTES = 80ULL * 1024ULL * 1024ULL;
cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, MAX_PERSIST_BYTES);
cudaStreamAttrValue attr = {};
attr.accessPolicyWindow.base_ptr = const_cast<void*>((const void*)w_base);
attr.accessPolicyWindow.num_bytes = std::min(w_bytes, MAX_PERSIST_BYTES);
attr.accessPolicyWindow.hitRatio = 1.0f;
attr.accessPolicyWindow.hitProp = cudaAccessPropertyPersisting;
attr.accessPolicyWindow.missProp = cudaAccessPropertyStreaming;
cudaStreamSetAttribute(stream, cudaStreamAttributeAccessPolicyWindow, &attr);
s_last_w_base = (const void*)w_base;
}
DecodeParams params;
const __nv_bfloat16* ptr = w_base;
for (int l = 0; l < num_layers; l++) {
params.layers[l].down_proj = ptr; ptr += HIDDEN_SIZE * INTERMEDIATE_SIZE;
params.layers[l].o_proj = ptr; ptr += HIDDEN_SIZE * Q_SIZE;
params.layers[l].input_ln = ptr; ptr += HIDDEN_SIZE;
params.layers[l].q_norm = ptr; ptr += HEAD_DIM;
params.layers[l].k_norm = ptr; ptr += HEAD_DIM;
params.layers[l].post_ln = ptr; ptr += HIDDEN_SIZE;
params.layers[l].qkv_proj = ptr; ptr += QKV_SIZE * HIDDEN_SIZE;
}
for (int l = 0; l < num_layers; l++) {
params.layers[l].gate_proj = ptr; ptr += INTERMEDIATE_SIZE * HIDDEN_SIZE;
params.layers[l].up_proj = ptr; ptr += INTERMEDIATE_SIZE * HIDDEN_SIZE;
params.layers[l].k_cache = reinterpret_cast<__nv_bfloat16*>(k_caches[l].data_ptr());
params.layers[l].v_cache = reinterpret_cast<__nv_bfloat16*>(v_caches[l].data_ptr());
}
float* ws = workspace_f32.data_ptr<float>();
float* g_residual = ws; ws += HIDDEN_SIZE;
float* g_activations = ws; ws += HIDDEN_SIZE;
float* g_qkv = ws; ws += QKV_SIZE;
float* g_attn_out = ws; ws += Q_SIZE;
float* g_mlp_intermediate = ws; ws += INTERMEDIATE_SIZE;
float* g_split_max = ws; ws += NUM_Q_HEADS * MAX_SPLITS_PER_KV;
float* g_split_sum = ws; ws += NUM_Q_HEADS * MAX_SPLITS_PER_KV;
float* g_split_acc = ws; ws += NUM_Q_HEADS * MAX_SPLITS_PER_KV * HEAD_DIM;
const __nv_bfloat16* x_rand_ptr = reinterpret_cast<const __nv_bfloat16*>(x_rand_stream.data_ptr());
const __nv_bfloat16* hidden_in_ptr = reinterpret_cast<const __nv_bfloat16*>(hidden_in.data_ptr());
__nv_bfloat16* hidden_out_ptr = reinterpret_cast<__nv_bfloat16*>(hidden_out.data_ptr());
unsigned int* locks_ptr = reinterpret_cast<unsigned int*>(attn_locks.data_ptr());
void* kernel_args[] = {
(void*)¶ms,
(void*)&x_rand_ptr,
(void*)&hidden_in_ptr,
(void*)&hidden_out_ptr,
(void*)&g_residual,
(void*)&g_activations,
(void*)&g_qkv,
(void*)&g_attn_out,
(void*)&g_mlp_intermediate,
(void*)&g_split_max,
(void*)&g_split_sum,
(void*)&g_split_acc,
(void*)&locks_ptr,
(void*)&num_layers,
(void*)&start_pos,
(void*)&n_steps,
(void*)&max_seq_len,
(void*)&attn_scale
};
cudaLaunchCooperativeKernel(
(void*)megaqwen_decode_kernel,
dim3(DEC_NUM_BLOCKS),
dim3(DEC_BLOCK_SIZE),
kernel_args,
0,
stream
);
}
"""
CPP_SRC = r"""
#include <torch/extension.h>
#include <vector>
void launch_megaqwen_decode(
torch::Tensor x_rand_stream,
torch::Tensor hidden_in,
torch::Tensor hidden_out,
torch::Tensor flat_weights,
std::vector<torch::Tensor> k_caches,
std::vector<torch::Tensor> v_caches,
torch::Tensor workspace_f32,
torch::Tensor attn_locks,
int num_layers,
int start_pos,
int n_steps,
int max_seq_len,
float attn_scale
);
"""
_ext = load_inline(
name="megaqwen_cuda_v7",
cpp_sources=[CPP_SRC],
cuda_sources=[CUDA_SRC],
functions=["launch_megaqwen_decode"],
extra_cuda_cflags=["-O3", "--use_fast_math", "-arch=sm_120"],
verbose=False,
)
class Block(nn.Module):
def __init__(self):
super().__init__()
H, I, D = HIDDEN, INTERMEDIATE, HEAD_DIM
self.input_ln = nn.Parameter(torch.ones(H, dtype=torch.bfloat16))
self.q_proj = nn.Parameter(torch.empty(NUM_Q * D, H, dtype=torch.bfloat16))
self.k_proj = nn.Parameter(torch.empty(NUM_KV * D, H, dtype=torch.bfloat16))
self.v_proj = nn.Parameter(torch.empty(NUM_KV * D, H, dtype=torch.bfloat16))
self.q_norm = nn.Parameter(torch.ones(D, dtype=torch.bfloat16))
self.k_norm = nn.Parameter(torch.ones(D, dtype=torch.bfloat16))
self.o_proj = nn.Parameter(torch.empty(H, NUM_Q * D, dtype=torch.bfloat16))
self.post_ln = nn.Parameter(torch.ones(H, dtype=torch.bfloat16))
self.gate_proj = nn.Parameter(torch.empty(I, H, dtype=torch.bfloat16))
self.up_proj = nn.Parameter(torch.empty(I, H, dtype=torch.bfloat16))
self.down_proj = nn.Parameter(torch.empty(H, I, dtype=torch.bfloat16))
for p in self.parameters():
if p is self.input_ln or p is self.post_ln or p is self.q_norm or p is self.k_norm:
continue
nn.init.normal_(p, std=0.02)
class Model(nn.Module):
def __init__(self, num_layers: int = NUM_LAYERS, max_seq: int = 131072):
super().__init__()
self.num_layers = num_layers
self.max_seq = max_seq
self.blocks = nn.ModuleList([Block() for _ in range(num_layers)])
self._flat_weights = None
self._pack_parameters_storage()
def _pack_parameters_storage(self):
device = self.blocks[0].input_ln.device
dtype = self.blocks[0].input_ln.dtype
if (
self._flat_weights is not None
and self._flat_weights.device == device
and self._flat_weights.dtype == dtype
and self.blocks[0].down_proj.data_ptr() == self._flat_weights.data_ptr()
):
return
total_elems = sum(p.numel() for p in self.parameters())
flat = torch.zeros(total_elems + 16384, dtype=dtype, device=device)
offset = 0
with torch.no_grad():
for blk in self.blocks:
for p in (
blk.down_proj,
blk.o_proj,
blk.input_ln,
blk.q_norm,
blk.k_norm,
blk.post_ln,
blk.q_proj,
blk.k_proj,
blk.v_proj,
):
n = p.numel()
flat[offset : offset + n].copy_(p.data.view(-1))
p.data = flat[offset : offset + n].view(p.shape)
offset += n
for blk in self.blocks:
for p in (blk.gate_proj, blk.up_proj):
n = p.numel()
flat[offset : offset + n].copy_(p.data.view(-1))
p.data = flat[offset : offset + n].view(p.shape)
offset += n
self._flat_weights = flat
def _apply(self, fn):
super()._apply(fn)
self._pack_parameters_storage()
return self
_WORKSPACES: dict[torch.device, tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]] = {}
def _get_workspace(device: torch.device):
ws = _WORKSPACES.get(device)
if ws is None:
ws_f32 = torch.zeros(128 * 1024, dtype=torch.float32, device=device)
locks = torch.zeros(32, dtype=torch.int32, device=device)
x_rand_pinned = torch.empty((256, HIDDEN), dtype=torch.bfloat16, device="cpu", pin_memory=True)
x_rand_gpu = torch.empty((256, HIDDEN), dtype=torch.bfloat16, device=device)
h_out = torch.empty(HIDDEN, dtype=torch.bfloat16, device=device)
ws = (ws_f32, locks, x_rand_pinned, x_rand_gpu, h_out)
_WORKSPACES[device] = ws
return ws
def empty_caches(num_layers: int, max_seq: int, device, dtype=torch.bfloat16):
k = [
torch.empty(NUM_KV, max_seq, HEAD_DIM, device=device, dtype=dtype)
for _ in range(num_layers)
]
v = [
torch.empty(NUM_KV, max_seq, HEAD_DIM, device=device, dtype=dtype)
for _ in range(num_layers)
]
return k, v
def _seeded_hidden(seed: int, device) -> torch.Tensor:
g = torch.Generator(device="cpu")
g.manual_seed(seed)
return torch.randn(HIDDEN, generator=g, dtype=torch.bfloat16).to(device)
_PREFILL_CACHE: dict = {}
@torch.no_grad()
def prefill(
model: Model,
ctx_len: int,
seed: int,
device: torch.device | None = None,
):
device = device or next(model.parameters()).device
if model._flat_weights is None or model._flat_weights.device != device:
model = model.to(device).eval()
assert ctx_len <= model.max_seq
w_ptr = int(model._flat_weights.data_ptr())
w_sample = (
float(model._flat_weights[0].item()),
float(model._flat_weights[HIDDEN].item()),
float(model._flat_weights[HIDDEN * 1024].item()),
)
cache_key = (id(model), w_ptr, w_sample, ctx_len, model.max_seq, seed)
if _PREFILL_CACHE.get("key") == cache_key:
h_saved, k_caches, v_caches = _PREFILL_CACHE["val"]
return h_saved.clone(), k_caches, v_caches
prev = _PREFILL_CACHE.get("buffers")
if (
prev is not None
and len(prev[0]) == model.num_layers
and prev[0][0].shape == (NUM_KV, model.max_seq, HEAD_DIM)
and prev[0][0].device == device
):
k_caches, v_caches = prev
else:
k_caches, v_caches = empty_caches(model.num_layers, model.max_seq, device)
_PREFILL_CACHE["buffers"] = (k_caches, v_caches)
h = _seeded_hidden(seed, device)
if ctx_len > 0:
g = torch.Generator(device="cpu")
g.manual_seed(seed + 1)
x_rand = torch.randn((ctx_len, HIDDEN), generator=g, dtype=torch.bfloat16).to(device)
ws_f32, locks, _, _, _ = _get_workspace(device)
_ext.launch_megaqwen_decode(
x_rand,
h,
h,
model._flat_weights,
k_caches,
v_caches,
ws_f32,
locks,
model.num_layers,
0,
ctx_len,
model.max_seq,
1.0 / math.sqrt(HEAD_DIM),
)
_PREFILL_CACHE["key"] = cache_key
_PREFILL_CACHE["val"] = (h.clone(), k_caches, v_caches)
return h, k_caches, v_caches
@torch.no_grad()
def decode_steps(
model: Model,
hidden: torch.Tensor,
k_caches: list[torch.Tensor],
v_caches: list[torch.Tensor],
start_pos: int,
n_steps: int,
seed: int,
):
device = hidden.device
if n_steps <= 0:
return hidden.clone(), k_caches, v_caches
ws_f32, locks, x_rand_pinned, x_rand_gpu, h_out = _get_workspace(device)
g = torch.Generator(device="cpu")
g.manual_seed(seed + 2)
if n_steps <= x_rand_pinned.shape[0]:
cpu_slice = x_rand_pinned[:n_steps]
torch.randn((n_steps, HIDDEN), generator=g, dtype=torch.bfloat16, out=cpu_slice)
x_rand = x_rand_gpu[:n_steps]
x_rand.copy_(cpu_slice, non_blocking=True)
else:
x_rand = torch.randn((n_steps, HIDDEN), generator=g, dtype=torch.bfloat16).to(device)
max_seq = k_caches[0].shape[1]
_ext.launch_megaqwen_decode(
x_rand,
hidden,
h_out,
model._flat_weights,
k_caches,
v_caches,
ws_f32,
locks,
model.num_layers,
start_pos,
n_steps,
max_seq,
1.0 / math.sqrt(HEAD_DIM),
)
return h_out, k_caches, v_caches
def run(
ctx_len: int,
n_decode: int,
seed: int,
model: Model | None = None,
max_seq: int | None = None,
) -> dict:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
max_seq = max_seq or max(ctx_len + n_decode, 512)
if model is None:
model = Model(NUM_LAYERS, max_seq)
else:
if getattr(model, "max_seq", 0) < ctx_len + n_decode:
raise ValueError(
f"model.max_seq={getattr(model, 'max_seq', None)} too small for "
f"ctx_len={ctx_len}+n_decode={n_decode}"
)
model = model.to(device).eval()
h, k_caches, v_caches = prefill(model, ctx_len, seed, device=device)
h, k_caches, v_caches = decode_steps(
model, h, k_caches, v_caches, start_pos=ctx_len, n_steps=n_decode, seed=seed
)
return {
"last_hidden": h.detach(),
"ctx_len": ctx_len,
"decode_steps": n_decode,
}
def get_init_inputs():
return [NUM_LAYERS, 131072]
def get_inputs():
return []
20261013_175315_agy_gemini-4-argon-high_03_megaqwen_decode