MegaQwen Decode GPT-6 Sol

"""CUDA decode for the four-block Qwen3-0.6B geometry."""
from __future__ import annotations

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

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


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.input_ln = nn.Parameter(torch.ones(1024, dtype=torch.bfloat16))
        self.q_proj = nn.Parameter(torch.empty(2048, 1024, dtype=torch.bfloat16))
        self.k_proj = nn.Parameter(torch.empty(1024, 1024, dtype=torch.bfloat16))
        self.v_proj = nn.Parameter(torch.empty(1024, 1024, dtype=torch.bfloat16))
        self.q_norm = nn.Parameter(torch.ones(128, dtype=torch.bfloat16))
        self.k_norm = nn.Parameter(torch.ones(128, dtype=torch.bfloat16))
        self.o_proj = nn.Parameter(torch.empty(1024, 2048, dtype=torch.bfloat16))
        self.post_ln = nn.Parameter(torch.ones(1024, dtype=torch.bfloat16))
        self.gate_proj = nn.Parameter(torch.empty(3072, 1024, dtype=torch.bfloat16))
        self.up_proj = nn.Parameter(torch.empty(3072, 1024, dtype=torch.bfloat16))
        self.down_proj = nn.Parameter(torch.empty(1024, 3072, dtype=torch.bfloat16))
        for p in self.parameters():
            if p.ndim == 2:
                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)])


_CUDA = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cmath>
using bf = __nv_bfloat16;

__device__ __forceinline__ float warp_sum(float v) {
  for (int d=16;d>0;d>>=1) v += __shfl_down_sync(0xffffffff,v,d);
  return v;
}
__device__ __forceinline__ float block_sum(float v) {
  __shared__ float s[8];
  int lane=threadIdx.x&31, warp=threadIdx.x>>5;
  v=warp_sum(v);
  if(lane==0) s[warp]=v;
  __syncthreads();
  v=(threadIdx.x<blockDim.x/32)?s[lane]:0.f;
  if(warp==0) v=warp_sum(v);
  return __shfl_sync(0xffffffff,v,0);
}

__global__ void norm_input(const bf* prev,const bf* rnd,const bf* w,
                           float* residual,float* norm,int mix) {
  int i=threadIdx.x;
  __shared__ float ss[8];
  float sum=0.f;
  for(int j=i;j<1024;j+=256) {
    float x=__bfloat162float(prev[j]);
    if(mix) {
      bf a=__float2bfloat16_rn(0.5f*__bfloat162float(rnd[j]));
      bf b=__float2bfloat16_rn(0.5f*x);
      x=__bfloat162float(__float2bfloat16_rn(__bfloat162float(a)+__bfloat162float(b)));
    }
    residual[j]=x;
    sum+=x*x;
  }
  sum=warp_sum(sum);
  if((i&31)==0) ss[i>>5]=sum;
  __syncthreads();
  float v=(i<8)?ss[i]:0.f;
  if(i<32) v=warp_sum(v);
  if(i==0) ss[0]=rsqrtf(v/1024.f+1e-6f);
  __syncthreads();
  float inv=ss[0];
  for(int j=i;j<1024;j+=256) norm[j]=residual[j]*inv*__bfloat162float(w[j]);
}

// One warp per output row, with eight bf16 weights per lane per iteration.
__global__ void qkv_mv(const float* x,const bf* qw,const bf* kw,const bf* vw,
                       float* q,float* k,float* v) {
  int row=blockIdx.x*4+(threadIdx.x>>5), lane=threadIdx.x&31;
  if(row>=4096) return;
  const bf* w; float* out;
  if(row<2048) {w=qw+row*1024;out=q+row;}
  else if(row<3072) {w=kw+(row-2048)*1024;out=k+row-2048;}
  else {w=vw+(row-3072)*1024;out=v+row-3072;}
  float s=0.f;
  #pragma unroll
  for(int base=lane*8;base<1024;base+=256) {
    #pragma unroll
    for(int j=0;j<8;j++) s=fmaf(__bfloat162float(w[base+j]),x[base+j],s);
  }
  s=warp_sum(s);
  if(lane==0) *out=s;
}

__global__ void qk_rope(const float* q,const float* k,const float* v,
                        const bf* qw,const bf* kw,const float* cosine,const float* sine,
                        float* qr,bf* kc,bf* vc,
                        int pos,int stride) {
  int head=blockIdx.x, d=threadIdx.x;
  __shared__ float vals[128], shared[4];
  bool isq=head<16;
  int h=isq?head:head-16;
  const float* src=(isq?q:k)+h*128;
  float x=src[d];
  vals[d]=x;
  float sq=x*x;
  sq=warp_sum(sq);
  if((d&31)==0) shared[d>>5]=sq;
  __syncthreads();
  if(d==0) shared[0]=rsqrtf((shared[0]+shared[1]+shared[2]+shared[3])/128.f+1e-6f);
  __syncthreads();
  float n=vals[d]*shared[0]*__bfloat162float((isq?qw:kw)[d]);
  vals[d]=n;
  __syncthreads();
  int f=d&63;
  float c=cosine[pos*64+f],s=sine[pos*64+f];
  float y=d<64 ? n*c-vals[d+64]*s : vals[d-64]*s+n*c;
  if(isq) qr[h*128+d]=y;
  else {
    kc[(h*stride+pos)*128+d]=__float2bfloat16_rn(y);
    vc[(h*stride+pos)*128+d]=__float2bfloat16_rn(v[h*128+d]);
  }
}

__global__ void scores_single(const float* q,const bf* kc,float* scores,int len,int stride) {
  int head=blockIdx.y,warp=threadIdx.x>>5,lane=threadIdx.x&31;
  int t=blockIdx.x*128+warp;
  for(int j=0;j<32;j++,t+=4) {
    if(t>=len) break;
    const bf* key=kc+((head>>1)*stride+t)*128;
    const float* query=q+head*128;
    float s=0.f;
    #pragma unroll
    for(int d=lane;d<128;d+=32) s=fmaf(query[d],__bfloat162float(key[d]),s);
    s=warp_sum(s);
    if(lane==0) scores[head*stride+t]=s*0.08838834764831845f;
  }
}

__global__ void partial_single(const float* scores,const bf* vc,float* partial,
                               float* stats,int len,int stride,int chunks) {
  int head=blockIdx.y,chunk=blockIdx.x,d=threadIdx.x;
  int begin=chunk*256,end=min(begin+256,len),n=end-begin;
  __shared__ float sm[256],weights[256],red[128];
  float a=-INFINITY,b=-INFINITY;
  if(d<n) a=scores[head*stride+begin+d];
  if(d+128<n) b=scores[head*stride+begin+d+128];
  sm[d]=fmaxf(a,b);
  __syncthreads();
  for(int off=64;off>0;off>>=1) {
    if(d<off) sm[d]=fmaxf(sm[d],sm[d+off]);
    __syncthreads();
  }
  float mx=sm[0];
  float e0=d<n?expf(a-mx):0.f;
  float e1=d+128<n?expf(b-mx):0.f;
  weights[d]=e0;weights[d+128]=e1;red[d]=e0+e1;
  __syncthreads();
  for(int off=64;off>0;off>>=1) {
    if(d<off) red[d]+=red[d+off];
    __syncthreads();
  }
  float acc=0.f;
  for(int t=0;t<n;t++)
    acc=fmaf(weights[t],__bfloat162float(vc[((head>>1)*stride+begin+t)*128+d]),acc);
  partial[(head*chunks+chunk)*128+d]=acc;
  if(d==0) {stats[(head*chunks+chunk)*2]=mx;stats[(head*chunks+chunk)*2+1]=red[0];}
}

__global__ void scores_kernel(const float* q,const bf* kc,float* scores,int len,int stride) {
  int kvhead=blockIdx.y, head=kvhead*2, warp=threadIdx.x>>5,lane=threadIdx.x&31;
  int t=blockIdx.x*128+warp;
  for(int j=0;j<32;j++,t+=4) {
    if(t>=len) break;
    const bf* key=kc+(kvhead*stride+t)*128;
    const float *query0=q+head*128,*query1=query0+128;
    float s0=0.f,s1=0.f;
    #pragma unroll
    for(int d=lane;d<128;d+=32) {
      float val=__bfloat162float(key[d]);
      s0=fmaf(query0[d],val,s0);
      s1=fmaf(query1[d],val,s1);
    }
    s0=warp_sum(s0);s1=warp_sum(s1);
    if(lane==0) {
      scores[head*stride+t]=s0*0.08838834764831845f;
      scores[(head+1)*stride+t]=s1*0.08838834764831845f;
    }
  }
}

__global__ void partial_attention(const float* scores,const bf* vc,float* partial,
                                  float* stats,int len,int stride,int chunks) {
  int kvhead=blockIdx.y,head=kvhead*2,chunk=blockIdx.x,d=threadIdx.x;
  int begin=chunk*256,end=min(begin+256,len),n=end-begin;
  __shared__ float sm[2][128], weights[2][256], red[2][128];
  float a0=-INFINITY,b0=-INFINITY,a1=-INFINITY,b1=-INFINITY;
  if(d<n) {
    a0=scores[head*stride+begin+d];
    a1=scores[(head+1)*stride+begin+d];
  }
  if(d+128<n) {
    b0=scores[head*stride+begin+d+128];
    b1=scores[(head+1)*stride+begin+d+128];
  }
  sm[0][d]=fmaxf(a0,b0);
  sm[1][d]=fmaxf(a1,b1);
  __syncthreads();
  for(int off=64;off>0;off>>=1) {
    if(d<off) {
      sm[0][d]=fmaxf(sm[0][d],sm[0][d+off]);
      sm[1][d]=fmaxf(sm[1][d],sm[1][d+off]);
    }
    __syncthreads();
  }
  float mx0=sm[0][0],mx1=sm[1][0];
  float e00=d<n?expf(a0-mx0):0.f;
  float e01=d+128<n?expf(b0-mx0):0.f;
  float e10=d<n?expf(a1-mx1):0.f;
  float e11=d+128<n?expf(b1-mx1):0.f;
  weights[0][d]=e00;
  weights[0][d+128]=e01;
  weights[1][d]=e10;
  weights[1][d+128]=e11;
  red[0][d]=e00+e01;
  red[1][d]=e10+e11;
  __syncthreads();
  for(int off=64;off>0;off>>=1) {
    if(d<off) {
      red[0][d]+=red[0][d+off];
      red[1][d]+=red[1][d+off];
    }
    __syncthreads();
  }
  float acc0=0.f,acc1=0.f;
  for(int t=0;t<n;t++) {
    float val=__bfloat162float(vc[(kvhead*stride+begin+t)*128+d]);
    acc0=fmaf(weights[0][t],val,acc0);
    acc1=fmaf(weights[1][t],val,acc1);
  }
  partial[(head*chunks+chunk)*128+d]=acc0;
  partial[((head+1)*chunks+chunk)*128+d]=acc1;
  if(d==0) {
    stats[(head*chunks+chunk)*2]=mx0;
    stats[(head*chunks+chunk)*2+1]=red[0][0];
    stats[((head+1)*chunks+chunk)*2]=mx1;
    stats[((head+1)*chunks+chunk)*2+1]=red[1][0];
  }
}

// One CTA handles a K/V tile for both query heads. It forms local scores,
// applies two softmaxes, and reads each V element once for both outputs.
__global__ void fused_attention(const float* q,const bf* kc,const bf* vc,
                                float* partial,float* stats,int len,int stride,int chunks) {
  int kvhead=blockIdx.y,head=kvhead*2,chunk=blockIdx.x;
  int tid=threadIdx.x,warp=tid>>5,lane=tid&31;
  int begin=chunk*256,n=min(256,len-begin);
  __shared__ float score0[256],score1[256],weights0[256],weights1[256];
  __shared__ float red0[256],red1[256],acc0[256],acc1[256];
  __shared__ float max0,max1,den0,den1;
  const float* q0=q+head*128;
  const float* q1=q0+128;
  for(int t=warp;t<n;t+=8) {
    const bf* key=kc+(kvhead*stride+begin+t)*128;
    float s0=0.f,s1=0.f;
    #pragma unroll
    for(int d=lane;d<128;d+=32) {
      float val=__bfloat162float(key[d]);
      s0=fmaf(q0[d],val,s0);
      s1=fmaf(q1[d],val,s1);
    }
    s0=warp_sum(s0);s1=warp_sum(s1);
    if(lane==0) {
      score0[t]=s0*0.08838834764831845f;
      score1[t]=s1*0.08838834764831845f;
    }
  }
  __syncthreads();
  red0[tid]=tid<n?score0[tid]:-INFINITY;
  red1[tid]=tid<n?score1[tid]:-INFINITY;
  __syncthreads();
  for(int off=128;off>0;off>>=1) {
    if(tid<off) {
      red0[tid]=fmaxf(red0[tid],red0[tid+off]);
      red1[tid]=fmaxf(red1[tid],red1[tid+off]);
    }
    __syncthreads();
  }
  if(tid==0) {max0=red0[0];max1=red1[0];}
  __syncthreads();
  float e0=tid<n?expf(score0[tid]-max0):0.f;
  float e1=tid<n?expf(score1[tid]-max1):0.f;
  weights0[tid]=e0;weights1[tid]=e1;
  red0[tid]=e0;red1[tid]=e1;
  __syncthreads();
  for(int off=128;off>0;off>>=1) {
    if(tid<off) {red0[tid]+=red0[tid+off];red1[tid]+=red1[tid+off];}
    __syncthreads();
  }
  if(tid==0) {den0=red0[0];den1=red1[0];}
  __syncthreads();
  int d=tid&127,half=tid>>7;
  float a0=0.f,a1=0.f;
  for(int t=half;t<n;t+=2) {
    float val=__bfloat162float(vc[(kvhead*stride+begin+t)*128+d]);
    a0=fmaf(weights0[t],val,a0);
    a1=fmaf(weights1[t],val,a1);
  }
  acc0[tid]=a0;acc1[tid]=a1;
  __syncthreads();
  if(tid<128) {
    partial[(head*chunks+chunk)*128+tid]=acc0[tid]+acc0[tid+128];
    partial[((head+1)*chunks+chunk)*128+tid]=acc1[tid]+acc1[tid+128];
  }
  if(tid==0) {
    stats[(head*chunks+chunk)*2]=max0;
    stats[(head*chunks+chunk)*2+1]=den0;
    stats[((head+1)*chunks+chunk)*2]=max1;
    stats[((head+1)*chunks+chunk)*2+1]=den1;
  }
}

__global__ void finish_attention(const float* partial,const float* stats,float* out,int chunks) {
  int head=blockIdx.x,d=threadIdx.x;
  __shared__ float mx,den,scale[1024];
  if(d==0) {
    float m=-INFINITY;
    for(int c=0;c<chunks;c++) m=fmaxf(m,stats[(head*chunks+c)*2]);
    mx=m;
  }
  __syncthreads();
  for(int c=d;c<chunks;c+=128) scale[c]=expf(stats[(head*chunks+c)*2]-mx);
  __syncthreads();
  if(d==0) {
    float s=0.f;
    for(int c=0;c<chunks;c++) s+=scale[c]*stats[(head*chunks+c)*2+1];
    den=s;
  }
  __syncthreads();
  float acc=0.f;
  for(int c=0;c<chunks;c++)
    acc=fmaf(scale[c],partial[(head*chunks+c)*128+d],acc);
  out[head*128+d]=acc/den;
}

__global__ void mv(const float* x,const bf* w,float* y,int rows,int cols) {
  int row=blockIdx.x*4+(threadIdx.x>>5),lane=threadIdx.x&31;
  if(row>=rows) return;
  const bf* wr=w+row*cols;
  float s=0.f;
  for(int base=lane*8;base<cols;base+=256) {
    #pragma unroll
    for(int j=0;j<8;j++) s=fmaf(__bfloat162float(wr[base+j]),x[base+j],s);
  }
  s=warp_sum(s);
  if(lane==0) y[row]=s;
}

__global__ void postnorm(const float* residual,const float* att,const bf* w,
                         float* postres,float* norm) {
  int i=threadIdx.x;
  __shared__ float ss[8];
  float sum=0.f;
  for(int j=i;j<1024;j+=256) {
    float x=residual[j]+att[j];
    postres[j]=x;
    sum+=x*x;
  }
  sum=warp_sum(sum);
  if((i&31)==0) ss[i>>5]=sum;
  __syncthreads();
  float v=i<8?ss[i]:0.f;
  if(i<32) v=warp_sum(v);
  if(i==0) ss[0]=rsqrtf(v/1024.f+1e-6f);
  __syncthreads();
  for(int j=i;j<1024;j+=256) norm[j]=postres[j]*ss[0]*__bfloat162float(w[j]);
}

__global__ void gate_up(const float* x,const bf* gw,const bf* uw,float* out) {
  int row=blockIdx.x*4+(threadIdx.x>>5),lane=threadIdx.x&31;
  if(row>=3072) return;
  const bf* g=gw+row*1024;const bf* u=uw+row*1024;
  float a=0.f,b=0.f;
  for(int base=lane*8;base<1024;base+=256) {
    #pragma unroll
    for(int j=0;j<8;j++) {
      float v=x[base+j];
      a=fmaf(__bfloat162float(g[base+j]),v,a);
      b=fmaf(__bfloat162float(u[base+j]),v,b);
    }
  }
  a=warp_sum(a);b=warp_sum(b);
  if(lane==0) out[row]=(a/(1.f+expf(-a)))*b;
}

__global__ void finish_layer(const float* postres,const float* down,bf* out) {
  int i=blockIdx.x*256+threadIdx.x;
  if(i<1024) out[i]=__float2bfloat16_rn(postres[i]+down[i]);
}

void layer(torch::Tensor prev,torch::Tensor rnd,torch::Tensor inw,
           torch::Tensor qw,torch::Tensor kw,torch::Tensor vw,
           torch::Tensor qnw,torch::Tensor knw,torch::Tensor ow,
           torch::Tensor postw,torch::Tensor gw,torch::Tensor uw,torch::Tensor dw,
           torch::Tensor cosine,torch::Tensor sine,
           torch::Tensor kc,torch::Tensor vc,torch::Tensor out,
           torch::Tensor scratch,torch::Tensor scores,torch::Tensor part,torch::Tensor stats,
           int pos,int mix) {
  auto stream=at::cuda::getCurrentCUDAStream();
  float* f=scratch.data_ptr<float>();
  // scratch: residual, norm, q, k, v, rotated q, attention, O result,
  // post residual, post norm, gate*up, down.
  float *res=f,*norm=f+1024,*q=f+2048,*k=f+4096,*v=f+5120;
  float *qr=f+6144,*att=f+8192,*o=f+10240,*post=f+11264,*pn=f+12288;
  float *gu=f+13312,*down=f+16384;
  int stride=kc.size(1),len=pos+1,chunks=(len+255)/256;
  norm_input<<<1,256,0,stream>>>(reinterpret_cast<const bf*>(prev.data_ptr()),
     reinterpret_cast<const bf*>(rnd.data_ptr()),reinterpret_cast<const bf*>(inw.data_ptr()),res,norm,mix);
  qkv_mv<<<1024,128,0,stream>>>(norm,reinterpret_cast<const bf*>(qw.data_ptr()),
    reinterpret_cast<const bf*>(kw.data_ptr()),reinterpret_cast<const bf*>(vw.data_ptr()),q,k,v);
  qk_rope<<<24,128,0,stream>>>(q,k,v,reinterpret_cast<const bf*>(qnw.data_ptr()),
    reinterpret_cast<const bf*>(knw.data_ptr()),cosine.data_ptr<float>(),sine.data_ptr<float>(),qr,
    reinterpret_cast<bf*>(kc.data_ptr()),reinterpret_cast<bf*>(vc.data_ptr()),pos,stride);
  if(len<16384) {
    scores_single<<<dim3((len+127)/128,16),128,0,stream>>>(qr,
      reinterpret_cast<const bf*>(kc.data_ptr()),scores.data_ptr<float>(),len,stride);
    partial_single<<<dim3(chunks,16),128,0,stream>>>(scores.data_ptr<float>(),
      reinterpret_cast<const bf*>(vc.data_ptr()),part.data_ptr<float>(),stats.data_ptr<float>(),len,stride,chunks);
  } else {
    fused_attention<<<dim3(chunks,8),256,0,stream>>>(qr,
      reinterpret_cast<const bf*>(kc.data_ptr()),reinterpret_cast<const bf*>(vc.data_ptr()),
      part.data_ptr<float>(),stats.data_ptr<float>(),len,stride,chunks);
  }
  finish_attention<<<16,128,0,stream>>>(part.data_ptr<float>(),stats.data_ptr<float>(),att,chunks);
  mv<<<256,128,0,stream>>>(att,reinterpret_cast<const bf*>(ow.data_ptr()),o,1024,2048);
  postnorm<<<1,256,0,stream>>>(res,o,reinterpret_cast<const bf*>(postw.data_ptr()),post,pn);
  gate_up<<<768,128,0,stream>>>(pn,reinterpret_cast<const bf*>(gw.data_ptr()),
    reinterpret_cast<const bf*>(uw.data_ptr()),gu);
  mv<<<256,128,0,stream>>>(gu,reinterpret_cast<const bf*>(dw.data_ptr()),down,1024,3072);
  finish_layer<<<4,256,0,stream>>>(post,down,reinterpret_cast<bf*>(out.data_ptr()));
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME,m) {m.def("layer",&layer);}
'''

_ext = None


def _extension():
    global _ext
    if _ext is None:
        _ext = load_inline(
            name="megaqwen_decode_sm120_v4",
            cpp_sources="",
            cuda_sources=_CUDA,
            extra_cuda_cflags=["-O3", "-arch=sm_120", "--use_fast_math"],
            with_cuda=True,
            verbose=False,
        )
    return _ext


class _Workspace:
    def __init__(self, model):
        device = next(model.parameters()).device
        if not hasattr(model, "_rope_cos") or model._rope_cos.shape[0] < model.max_seq:
            inv = 1.0 / (10000 ** (torch.arange(64, device=device, dtype=torch.float32) / 64))
            angles = torch.outer(torch.arange(model.max_seq, device=device, dtype=torch.float32), inv)
            model._rope_cos = angles.cos().contiguous()
            model._rope_sin = angles.sin().contiguous()
        self.scratch = torch.empty(17408, device=device, dtype=torch.float32)
        self.scores = torch.empty(16, model.max_seq, device=device, dtype=torch.float32)
        chunks = (model.max_seq + 255) // 256
        self.part = torch.empty(16, chunks, 128, device=device, dtype=torch.float32)
        self.stats = torch.empty(16, chunks, 2, device=device, dtype=torch.float32)
        self.next = torch.empty(1024, device=device, dtype=torch.bfloat16)
        self.other = torch.empty_like(self.next)


@torch.no_grad()
def _step(model, h, rnd, k_caches, v_caches, pos, work):
    ext = _extension()
    for i, block in enumerate(model.blocks):
        out = work.next if i % 2 == 0 else work.other
        ext.layer(
            h, rnd, block.input_ln, block.q_proj, block.k_proj, block.v_proj,
            block.q_norm, block.k_norm, block.o_proj, block.post_ln,
            block.gate_proj, block.up_proj, block.down_proj,
            model._rope_cos, model._rope_sin,
            k_caches[i], v_caches[i], out, work.scratch, work.scores,
            work.part, work.stats, pos, int(i == 0),
        )
        h = out
    return h


def _cpu_random(seed, steps, device):
    g = torch.Generator(device="cpu")
    g.manual_seed(seed)
    for _ in range(steps):
        yield torch.randn(1024, generator=g, dtype=torch.bfloat16).to(device)


@torch.no_grad()
def prefill(model, ctx_len, seed, device=None):
    device = device or next(model.parameters()).device
    model = model.to(device).eval()
    if ctx_len > model.max_seq:
        raise ValueError("context exceeds cache capacity")
    work = _Workspace(model)
    g = torch.Generator(device="cpu")
    g.manual_seed(seed)
    h = torch.randn(1024, generator=g, dtype=torch.bfloat16).to(device)
    kc = [torch.empty(8, model.max_seq, 128, device=device, dtype=torch.bfloat16) for _ in model.blocks]
    vc = [torch.empty_like(t) for t in kc]
    for pos, rnd in enumerate(_cpu_random(seed + 1, ctx_len, device)):
        h = _step(model, h, rnd, kc, vc, pos, work)
    return h.clone(), kc, vc


@torch.no_grad()
def decode_steps(model, hidden, k_caches, v_caches, start_pos, n_steps, seed):
    work = _Workspace(model)
    h = hidden
    for i, rnd in enumerate(_cpu_random(seed + 2, n_steps, hidden.device)):
        h = _step(model, h, rnd, k_caches, v_caches, start_pos + i, work)
    return h.clone(), k_caches, v_caches


_decode_steps_impl = decode_steps


def run(ctx_len, decode_steps, seed, model=None, max_seq=None):
    device = torch.device("cuda:0")
    max_seq = max_seq or max(ctx_len + decode_steps, 512)
    if model is None:
        model = Model(NUM_LAYERS, max_seq)
    elif model.max_seq < ctx_len + decode_steps:
        raise ValueError("model cache too short")
    model = model.to(device).eval()
    h, kc, vc = prefill(model, ctx_len, seed, device=device)
    h, kc, vc = _decode_steps_impl(model, h, kc, vc, ctx_len, decode_steps, seed)
    return {"last_hidden": h.detach(), "ctx_len": ctx_len, "decode_steps": decode_steps}