/* PSUEDOCODE */
void generate(Model& model, std::string prompt, int steps) {  std::vector<int> encoded = tokenizer.encode(prompt);  InferenceState s(model);    // 1. Prefill step: Forward the model on each prompt token, discarding   // the output. This lets the model read the prompt and hydrates the KV   // cache.  for (int token : encoded) {    model.forward(s, token);  }  // 2. Decode step: Forward the model repeatedly, generating 1 token at a time.  for (int i = 0; i < steps; i++) {    model.forward(s, encoded.back());    int next_token = sampler.sample(s.logits);    encoded.push_back(next_token);    std::cout << tokenizer.decode_one(next_token) << std::flush;    if (next_token == tokenizer.EOS) {      break;    }  }}        

/* PSUEDOCODE */
// InferenceState is the minimum set of buffers needed to// hold state during the forward pass and exists to avoid// extra allocationsvoid Model::forward(InferenceState& s, int token) {  // The embedding table maps token IDs to embedding vectors,  // which are copied into a buffer of the inference state  s.x = copy_embedding(token, this->token_embedding_table);    // Models consist of a sequence of transformer blocks which  // mutate the inference state in order  for (Block& b : this->blocks) {    b->block(s);  }    // Usually there is a layer norm right before the final classifier  s.x = layernorm(s.x, this->lm_head_prenorm_weights);  // Typically we end with a linear transform from (dim) -> (vocab_size)  s.logits = linear(s.x, this->lm_head_classifier_weights); }
void Block::block(InferenceState& s) {  s.x_resid = layernorm(s.x, this->att_prenorm_weights);  // Multi-head attention typically includes:   // 1. RoPE on input (element-wise mutation w/ sines/cosines)  // 2. QKV matmuls and updating the KV cache  // 3. Causal self-attention, softmax, and value mixing  // 4. Projection back into the residual stream  s.x_resid = multi_head_attn(    s.x_resid,    this->wq,     this->wk,     this->wv,     this->key_cache,    this->value_cache  );  s.x += s.x_resid;  s.x_resid = layernorm(s.x, this->ffn_prenorm_weights);  // On modern architectures like Llama, this is a GLU feedforward   // with 3 linear transforms, not a simple MLP:  // -> w2(F.silu(w1(x)) * w3(x))  // Some architectures also split the FFN into a mixture of experts.  s.x_resid = ffn(s.x_resid, this->w1, this->w2, this->w3);  s.x += s.x_resid;}         

static void matmul(float* xout, float* x, float* w, int n, int d) {  // W (d,n) @ x (n,) -> xout (d,)  int i;#pragma omp parallel for private(i)  for (i = 0; i < d; i++) {    float val = 0.0f;    for (int j = 0; j < n; j++) {      val += w[i * n + j] * x[j];    }    xout[i] = val;  }}        

// F16C code technically operates on 16-bit unsigned short integerstypedef uint16_t f16_t;
// matmul supporting float16 weights via the F16C extension, which allows// conversion into float32 values before calculations.static void matmul(float* xout, float* x, f16_t* w, int n, int d) {#if defined(__AVX2__) && defined(__F16C__)  // W (d,n) @ x (n,) -> xout (d,)  assert(n % 16 == 0);  int i;#pragma omp parallel for private(i)  for (i = 0; i < d; i++) {    // Vectorized dot product of w[i][:] and x[:] where w is a packed float16 array.    __m256 sumlo = _mm256_setzero_ps();    __m256 sumhi = _mm256_setzero_ps();    for (int j = 0; j < n; j+=16) {      // Extract the next set of 16 float16 weights from `w` and store them      // to two separate float32 vectors of width 8 (`wveclo_ps`, `wvechi_ps`)      __m256i wvec = _mm256_loadu_si256((__m256i*)&w[i * n + j]);      __m128i wveclo = _mm256_extractf128_si256(wvec, 0);      __m128i wvechi = _mm256_extractf128_si256(wvec, 1);      __m256 wveclo_ps = _mm256_cvtph_ps(wveclo);      __m256 wvechi_ps = _mm256_cvtph_ps(wvechi);      // Extract the next two float32 vectors of width 8 `xveclo`, `xvechi` from `x`      __m256 xveclo = _mm256_loadu_ps(&x[j]);      __m256 xvechi = _mm256_loadu_ps(&x[j + 8]);      // Compute vectorized FMAs: sumlo += wveclo * xveclo, sumhi += wvechi * xvechi      sumlo = _mm256_fmadd_ps(wveclo_ps, xveclo, sumlo);      sumhi = _mm256_fmadd_ps(wvechi_ps, xvechi, sumhi);    }    // Horizontally reduce width-8 float32 vectors sumlo, sumhi to a scalar.    __m256 sum8 = _mm256_add_ps(sumlo, sumhi);              // sum8[0:8] = sumlo[0:8] + sumhi[0:8]    __m128 sum4 = _mm_add_ps(                               // sum4[0:4] = sum8[0:4] + sum8[4:8]      _mm256_extractf128_ps(sum8, 0),       _mm256_extractf128_ps(sum8, 1)    );    __m128 sum1 = _mm_dp_ps(sum4, _mm_set1_ps(1.0f), 0xf1); // sum1[0] = dot(sum4, [1,1,1,1])    xout[i] = _mm_cvtss_f32(sum1);  }#else  assert(false && "float16 not supported on this platform");#endif}        

// mix self.w2(F.silu(self.w1(x)) * self.w3(x))// Note this is a feedforward with a GLU, not a simple MLP.matmul<<<c.hidden_dim, WARP_SIZE>>>(  w1(), s.xb(), c.dim, c.hidden_dim, s.hb());matmul<<<c.hidden_dim, WARP_SIZE>>>(  w3(), s.xb(), c.dim, c.hidden_dim, s.hb2());glu_gelu<<<  (c.hidden_dim + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK,   MAX_THREADS_PER_BLOCK,>>>(  s.hb(), s.hb2(), s.hb());matmul<<<c.dim, WARP_SIZE>>>(  w2(), s.hb(), c.hidden_dim, c.dim, s.xb2());// ffn residual back into xadd_residuals<<<  (c.dim + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK,  MAX_THREADS_PER_BLOCK>>>(  s.x(), s.xb2(), c.dim, s.x());        

// mix self.w2(F.silu(self.w1(x)) * self.w3(x))// Note this is a feedforward with a GLU, not a simple MLP.matmul(s.hb(), s.xb(), w1<T>(), c.dim, c.hidden_dim);matmul(s.hb2(), s.xb(), w3<T>(), c.dim, c.hidden_dim);for (int i = 0; i < c.hidden_dim; ++i) {  s.hb()[i] = gelu(s.hb()[i]) * s.hb2()[i];}matmul(s.xb2(), s.hb(), w2<T>(), c.hidden_dim, c.dim);// residual connection back into xfor (int i = 0; i < c.dim; ++i) {  s.x()[i] += s.xb2()[i];}        

__global__void matmul(const float* A, const float* x, int n, int d, float* out) {  // A (d,n) @ x (n,) -> out (d,)  int i = blockIdx.x * blockDim.x + threadIdx.x;  if (i >= d) return;  float sum = 0.0;  for (int j = 0; j < n; j++) {    sum += A[n * i + j] * x[j];  }  out[i] = sum;}
/* usage */int MAX_THREADS_PER_BLOCK = 1024;matmul<<<  (d + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK,   MAX_THREADS_PER_BLOCK>>>(A, x, n, d, out);

__device__ inline float warp_reduce_sum(float val) {  for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)    val += __shfl_down_sync(0xffffffff, val, offset);
  return val;}
__device__inline float matmul_row(const float* row, const float* x, int offset, int dim) {  float sum = 0.0;  for (int j = offset; j < dim; j += WARP_SIZE) {    float v = row[j] * x[j];    sum += v;  }  return warp_reduce_sum(sum);}
__global__void matmul(const float* A, const float* x, int n, int d, float* out) {  // A (d,n) @ x (n,) -> out (d,)  // PRECOND: Blocks are 1-D and same size as warp.  int i = blockIdx.x;  if (i >= d) return;  int offset = threadIdx.x;  float rowSum = matmul_row(&A[n * i], x, offset, n);  if (threadIdx.x == 0) {    out[i] = rowSum;  }}
/* usage */int BLOCK_SIZE = WARP_SIZE;matmul<<<d, BLOCK_SIZE>>>(A, x, n, d, out);        

Section: Memory Workload Analysis    --------------------------- ------------ ------------    Metric Name                  Metric Unit Metric Value    --------------------------- ------------ ------------    Memory Throughput           Gbyte/second       533.22    Mem Busy                               %        24.82    Max Bandwidth                          %        90.26    L1/TEX Hit Rate                        %        65.94    L2 Compression Success Rate            %            0    L2 Compression Ratio                                0    L2 Hit Rate                            %         2.03    Mem Pipes Busy                         %        28.33    --------------------------- ------------ ------------
    Section: Memory Workload Analysis Tables    ...    ----- --------------------------------------------------------------------------------------------------------------    WRN   The memory access pattern for stores from L1TEX to L2 is not optimal. The granularity of an L1TEX request to            L2 is a 128 byte cache line. That is 4 consecutive 32-byte sectors per L2 request. However, this kernel only            accesses an average of 1.0 sectors out of the possible 4 sectors per cache line. Check the Source Counters              section for uncoalesced stores and try to minimize how many cache lines need to be accessed per memory                  request.                                                                                                          ----- --------------------------------------------------------------------------------------------------------------    ...        

__device__ inline float blocktranspose(float v, float def) {  // Performs block-and-warp transpose operation:  //   For a block containing K warps where lane 0 contains val_k,  //   this function returns:  //   - For warp 0, lane K: val_k  //   - For all other warps and lanes: def  int lane = threadIdx.x % warpSize;  int warp = threadIdx.x / warpSize;    // Will hold results of all warps.  // Each lane of the warp accumulates across 1 head element at a time.  // NOTE: Assumes warpSize is 32  __shared__ float sm[32];  if (lane == 0) sm[warp] = v;  __syncthreads();    return lane < blockDim.x / warpSize ? sm[lane] : def;}
template <typename T>__global__void matmul_wide(const T* A, const float* x, int n, int d, float* out) {  // A (d,n) @ x (n,) -> out (d,)  // PRECOND: Block is 1-D and contains WPB warps.  int i = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;  if (i >= d) return;  // Warp j computes sum for row at <blockIdx.x*WPB + j>  // Lane 0 of each warp will hold result  int k = threadIdx.x % warpSize;  float rowSum = matmul_row(&A[n * i], x, k, n);  // Transpose values so lane k in warp 0 contains row at <blockIdx.x*WPB + k>  // For WPB=32, this allows us to coalesce 32 float32 writes into a single 128-byte store  rowSum = blocktranspose(rowSum, 1.0);  if (threadIdx.x < blockDim.x / warpSize) {    int block_start_i = blockIdx.x * blockDim.x / warpSize;    out[block_start_i + k] = rowSum;  }}

att_mix(const float *, const float *, int, int, int, int, int, float *) (32, 128, 1)x(32, 1, 1), Context 1, Stream 7, Device 0, CC 8.6    Section: GPU Speed Of Light Throughput    ----------------------- ------------- ------------    Metric Name               Metric Unit Metric Value    ----------------------- ------------- ------------    DRAM Frequency          cycle/nsecond         6.27    SM Frequency            cycle/nsecond         1.33    Elapsed Cycles                  cycle      5765045    Memory Throughput                   %         8.81    DRAM Throughput                     %         1.68    Duration                      msecond         4.35    L1/TEX Cache Throughput             %         5.17    L2 Cache Throughput                 %         8.81    SM Active Cycles                cycle   5685841.81    Compute (SM) Throughput             %         0.47    ----------------------- ------------- ------------
    WRN   This kernel exhibits low compute throughput and memory bandwidth utilization relative to the peak performance           of this device. Achieved compute throughput and/or memory bandwidth below 60.0% of peak typically indicate              latency issues. Look at Scheduler Statistics and Warp State Statistics for potential reasons.      

__device__ inline float warp_reduce_sum(float val) {for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)    val += __shfl_down_sync(0xffffffff, val, offset);return val;}
__global__void att_mix(  const float* vb, // (max_seq_len, n_kv_heads, head_dim)   const float* att, // (n_heads, kv_len)  int head_dim,   int n_heads,   int n_kv_heads,  int seq_len,   int max_seq_len,   float* out // (n_heads, head_dim)) {// PRECOND: blocks are 1-D and blockDim.x == WARP_SIZE  int h = blockIdx.x;  int group_size = n_heads / n_kv_heads;  int g = h / group_size;  int i = blockIdx.y;  int offset = threadIdx.x;  int kv_stride = n_kv_heads * head_dim;    const float* atth = att + max_seq_len * h;  const float* vh = vb + head_dim * g;  float* outh = out + head_dim * h;    float sum = 0.0;for (int t = offset; t < seq_len; t += WARP_SIZE) {    sum += vh[kv_stride * t + i] * atth[t];}  sum = warp_reduce_sum(sum);if (offset == 0) outh[i] = sum;}/* usage */dim3 tpb;tpb.x = WARP_SIZE;dim3 blocks;blocks.x = n_heads;blocks.y = head_dim;att_mix<<<blocks, tpb>>>(  vb, att,   head_dim, n_heads, n_kv_heads,   seq_len, max_seq_len, out);

      __global__void att_mix(  const float* vb,  // (max_seq_len, n_kv_heads, head_dim)   const float* att, // (n_heads, kv_len)  int head_dim,   int n_heads,   int n_kv_heads,  int seq_len,   int max_seq_len,   float* out // (n_heads, head_dim)) {// PRECOND: blocks are 1-D and `out` has been zeroed  int h = blockIdx.x;  int group_size = n_heads / n_kv_heads;  int g = h / group_size;  int kv_stride = n_kv_heads * head_dim;    const float* atth = att + max_seq_len * h;  const float* vh = vb + head_dim * g;  float* outh = out + head_dim * h;    int t_per_thread = seq_len / gridDim.y;  int t_start = blockIdx.y * t_per_thread;for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {    float sum = 0.0;for (int t = t_start; t < t_start + t_per_thread; t++) {      sum += vh[kv_stride * t + i] * atth[t];            }atomicAdd(&outh[i], sum);}}
int max_t_per_thread = 256;
dim3 tpb;tpb.x = warp_size;dim3 blocks;blocks.x = c.n_heads;blocks.y = (kv_len + max_t_per_thread - 1) / max_t_per_thread;cudaMemset(s.xb2(), 0, c.n_heads * c.head_dim * sizeof(float));att_mix<<<blocks, tpb>>>(  vb, s.att(),  c.head_dim, c.n_heads, c.n_kv_heads,   kv_len, c.max_seq_len, s.xb2());

__global__void att_mix(  const float* vb,  // (max_seq_len, n_kv_heads, head_dim)   const float* att, // (n_heads, kv_len)  int head_dim,   int n_heads,   int n_kv_heads,  int seq_len,   int max_seq_len,   float* out // (n_heads, head_dim)) {  // PRECOND: blocks are 2-D (warp_size, t_stride)  int h = blockIdx.x;  int group_size = n_heads / n_kv_heads;  int g = h / group_size;  int kv_stride = n_kv_heads * head_dim;    const float* atth = att + max_seq_len * h;  const float* vh = vb + head_dim * g;  float* outh = out + head_dim * h;    int warp_id = threadIdx.y;  int t_stride = blockDim.y;    // Capacity 32 since there can be at most 32 warps in a block.  __shared__ float shared[32];    for (int i = threadIdx.x; i < head_dim; i += warpSize) {    if (warp_id == 0) {      shared[threadIdx.x] = 0;    }    __syncthreads();    float sum = 0.0;    for (int t = warp_id; t < seq_len; t += t_stride) {      sum += vh[kv_stride * t + i] * atth[t];            }    atomicAdd(&shared[threadIdx.x], sum);    __syncthreads();    if (warp_id == 0) {      outh[i] = shared[threadIdx.x];      shared[threadIdx.x] = 0;    }  }}
dim3 tpb;tpb.x = warp_size;tpb.y = min(kv_len, max_threads_per_block / warp_size);dim3 blocks;blocks.x = c.n_heads;att_mix<<<blocks, tpb>>>(  vb, s.att(),  c.head_dim, c.n_heads, c.n_kv_heads,   kv_len, c.max_seq_len, s.xb2());      

float sum = 0.0;for (int t = warp_id; t < seq_len; t += t_stride) {  sum += vh[kv_stride * t + i] * atth[t];        }// atomicAdd `sum` when we're done     

float sum = 0.0;for (int t = warp_id; t < seq_len; t += t_stride) {  sum += __half2float(vh[kv_stride * t + i]) * atth[t];        }// atomicAdd `sum` when we're done        

float2 sum01 = make_float2(0.0, 0.0);for (int t = warp_id; t < seq_len; t += t_stride) {  float2 v01 = __half22float2(*((half2*)&vh[kv_stride * t + i]));  float att_t = atth[t];  // Sadly CUDA does not have float2 SIMD ops  sum01.x += v01.x * att_t;  sum01.y += v01.y * att_t;}// atomicAdd both `sum01` lanes when we're done

       float2 sum01 = make_float2(0.0, 0.0);constexpr int UNROLL = 16;half2 v01_0; float att_0; half2 v01_1; float att_1; half2 v01_2; float att_2; /* ... SNIP ... */half2 v01_15; float att_15;int t = warp_id;for (int ctr = 0; ctr < seq_len / t_stride; t += t_stride, ctr++) {  int ctr_mod = ctr % UNROLL;if (ctr_mod == 0) {// prefetch every UNROLL iterations    #define PREFETCH(j) \      v01_##j = *((half2*)&vh[kv_stride * (t + j*t_stride) + i]); \      att_##j = atth[t + j*t_stride];PREFETCH(0)PREFETCH(1)PREFETCH(2)/* ... SNIP ... */PREFETCH(15)    #undef PREFETCH}// pull one value out of prefetch batch  float2 v01;  float att_t;  switch (ctr_mod) {    #define CASE(j) \      case j: v01 = __half22float2(v01_##j); att_t = att_##j; break;CASE(0)CASE(1)CASE(2)/* ... SNIP ... */CASE(15)    #undef CASE}// Sadly CUDA does not have float2 SIMD ops  sum01.x += v01.x * att_t;  sum01.y += v01.y * att_t;}// Handle any loop remainder that can't be unrolledfor (; t < seq_len; t += t_stride) {  float2 v01 = __half22float2(*((half2*)&vh[kv_stride * t + i]));  float att_t = atth[t];  sum01.x += v01.x * att_t;  sum01.y += v01.y * att_t;}// atomicAdd both `sum01` lanes when we're done

import time
from transformers import AutoModelForCausalLM, AutoTokenizerimport torch
mistral_models_path = "../Mistral-7B-Instruct-v0.2"
tokenizer = AutoTokenizer.from_pretrained(mistral_models_path)
model = AutoModelForCausalLM.from_pretrained(mistral_models_path, torch_dtype=torch.float16)model.to("cuda")# short promptmodel_inputs = tokenizer(["Q: What is the meaning of life?"], return_tensors="pt").to("cuda")# do a warmup inferencegenerated_ids = model.generate(**model_inputs, max_new_tokens=1, do_sample=True)[0].tolist()# benchmarkstart_time = time.time() generated_ids = model.generate(**model_inputs, max_new_tokens=1000, do_sample=True)[0].tolist()# decode with mistral tokenizerresult = tokenizer.decode(generated_ids)end_time = time.time()
elapsed_s = end_time - start_timeprint(result)
num_prompt_tokens = model_inputs["input_ids"].shape[-1]num_generated_tokens = len(generated_ids) - num_prompt_tokens
print(f"Generation stats:\n" +      f"  prompt: {num_prompt_tokens} tokens\n" +      f"  generated: {num_generated_tokens} tokens\n" +      f"  throughput: {num_generated_tokens/elapsed_s}tok/s\n" +      f"  latency: {elapsed_s/num_generated_tokens}s/tok\n" +      f"  total: {elapsed_s}s\n")