PageSourceSearch

https://andrewkchan.dev/posts/yalm.html

html andrewkchan.dev collected 2026-10-03 23:08:47 UTC 74,914 bytes, 1,531 lines download raw bytes

1<!DOCTYPE html><html><head>
2  <title>⭐️ Fast LLM Inference From Scratch</title>
3  <meta charset="utf-8">
4  <!-- Google tag (gtag.js) -->
vendor: 69 bytes, lines 4-5
4
5  <script async="" src="https://www.googletagmanager.com/gtag/js?id=
5G-HYB0C59DXR
vendor: 14 bytes, lines 5-6
5"></script>
6  
6<!-- Sentry -->
7  
7<script src="https://js.sentry-cdn.com/f27e96506d16307fa97dcc9442b50117.min.js" crossorigin="anonymous"></script>
7
8  
8<script src="./template.v1.js"></script>
8
9  
9<script>
10    
vendor: 140 bytes, lines 10-14
10window.dataLayer = window.dataLayer || [];
11    function gtag(){dataLayer.push(arguments);}
12    gtag('js', new Date());
13
14    gtag('config', '
14G-HYB0C59DXR
vendor: 6 bytes, lines 14-15
14');
15  
15</script>
15
16  <style>
17    .toyb-nav ul {
18      padding-left: 0;
19    }
20    .toyb-nav li {
21      list-style-type: none;
22      margin-left: 0;
23    }
24    .toyb-nav-li-date {
25      color: #555;
26      font-size: 0.9em;
27      margin-right: 0.5em;
28    }
29  </style>
30
31    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/katex.min.css" integrity="sha384-9eLZqc9ds8eNjO3TmqPeYcDj8n+Qfa4nuSiGYa6DjLNcv9BtN69ZIulL9+8CqC9Y" crossorigin="anonymous">
32    
32<script defer="" src="https://cdn.jsdelivr.net/npm/[email protected]/dist/katex.min.js" integrity="sha384-K3vbOmF2BtaVai+Qk37uypf7VrgBubhQreNQe9aGsz9lB63dIFiQVlJbr92dw2Lx" crossorigin="anonymous"></script>
32
33
34    <!-- Google tag (gtag.js) -->
vendor: 71 bytes, lines 34-35
34
35    <script async="" src="https://www.googletagmanager.com/gtag/js?id=
35G-HYB0C59DXR
vendor: 16 bytes, lines 35-36
35"></script>
36    
36<script>
37      
vendor: 146 bytes, lines 37-41
37window.dataLayer = window.dataLayer || [];
38      function gtag(){dataLayer.push(arguments);}
39      gtag('js', new Date());
40
41      gtag('config', '
41G-HYB0C59DXR
vendor: 8 bytes, lines 41-42
41');
42    
42</script>
42
43
44    
44<script type="text/front-matter">
45      title: "Fast LLM Inference From Scratch"
46      published: December 12, 2024
47      description: "Pushing single-GPU inference throughput to the edge without libraries"
48      authors:
49      - Andrew Chan: http://andrewkchan.github.io
50      affiliations:
51      -
52    </script>
52
53
54    <style>
55      .compact-li li {
56        margin-bottom: 0.5em;
57      }
58      .caption {
59        font-size: 15px;
60        line-height: 1.3em;
61      }
62      .slider {
63          width: 100%;
64      }
65      #trex-forward-container {
66          width: min(90vw, 650px);
67      }
68      #trex-forward-main {
69          white-space: nowrap;
70      }
71      #trex-forward-slider-container {
72          width: 90%;
73          margin: 0 auto;
74      }
75      .trex-forward-carousel {
76          white-space: nowrap;
77      }
78      .trex-forward-carousel-img {
79          display: inline-block;
80          width: 19%;
81          opacity: 0.6;
82      }
83      .trex-forward-carousel-img:hover {
84          cursor: pointer;
85          opacity: 1.0;
86      }
87      #normalizing-flow>#normalizing-flow-gif {
88          display: none;
89      }
90      #normalizing-flow:hover>#normalizing-flow-gif {
91          display: inline-block;
92      }
93      #normalizing-flow:hover>#normalizing-flow-png {
94          display: none;
95      }
96      #toc {
97          line-height: 0.3em;
98          width: calc(984px - 648px);
99      }
100      @media (min-width: 1280px) {
101          #toc {
102              margin-right: 72px;
103              position: sticky; 
104              top: 72px;
105              z-index: 2;
106          }
107          #trex-noise-schedule-container { 
108              white-space: nowrap;
109          }
110          #trex-noise-schedule-gutter { 
111              margin-right:calc((100vw - 984px) / 2); 
112              z-index: 10 
113          }
114      }
115      #toc > li {
116          list-style-type: none;
117      }
118      .toc-anchor {
119        font-style: italic;
120      }
121      .l-gutter {
122          position: relative;
123          background-color: white;
124      }
125      #trex-noise-schedule-container {
126          width: min(90vw, 800px);
127      }
128    </style>
129  </head>
130<body>
131  <dt-article>
132    <h1><a href="/" class="hero">Andrew Chan</a></h1>
133  </dt-article>
134  <hr>
135  <article class="toyb-article">
136  
137  
138  
139  <dt-article>
140    <div class="l-gutter caption" id="toc"><h4>Contents</h4></div>
141    <h1>Fast LLM Inference From Scratch</h1>
142    <h2>Pushing single-GPU inference throughput to the edge without libraries</h2>
143    <dt-byline></dt-byline>
144    <ul class="caption">
145      <li><i>Source code for this article on <a href="https://github.com/andrewkchan/yalm">GitHub</a>.</i></li>
146      <li><i>Discussion on <a href="https://news.ycombinator.com/item?id=42417857">Hacker News</a>.</i></li>
147    </ul>
148    <p>
149      This post is about building an LLM inference engine using C++ and CUDA from scratch without libraries.
150    </p>
151    <p>
152      Why? In doing so, we can learn about the full stack of LLM inference - which is becoming increasingly important<dt-fn>Especially as inference compute becomes a new axis with which AI models scale,
153      and models are increasingly deployed locally to devices on the edge.</dt-fn> - from CUDA kernels to model architecture, and get a real sense of how different optimizations affect inference speed.
154      And one of the most important use cases is <i>running fast on a single prompt on consumer devices</i>.
155    </p>
156    <p>
157      That's what we'll focus on: building a program that can load weights of common open models and do single-batch inference on them on a single CPU + GPU server, and iteratively 
158      improving the token throughput until it surpasses <a href="https://github.com/ggerganov/llama.cpp">llama.cpp</a>. Readers should have basic familiarity with large language 
159      models, attention, and transformers. The full source code is available on GitHub: <a href="https://github.com/andrewkchan/yalm">yalm (Yet Another Language Model)</a>.
160    </p>
161    <h3>Acknowledgements</h3>
162    <ul class="compact-li">
163      <li>
164        <a href="https://github.com/zeux/calm">calm</a> - Much of my implementation is inspired by Arseny Kapoulkine's inference engine.
165        In a way, this project was kicked off by “understand calm and what makes it so fast.” I've tried to keep my code more readable for myself though, and as much as possible 
166        scientifically understanding optimizations, which means foregoing some advanced techniques used in calm like dynamic parallelism.
167      </li>
168      <li>
169        <a href="https://github.com/karpathy/llama2.c">llama2.c</a> - Parts of the CPU backend come from Andrej Karpathy's excellent C implementation of Llama inference.
170      </li>
171    </ul>
172    <dt-byline></dt-byline>
173    <h1 id="section-1">1. Recap: LLM architectures and inference</h1>
174    <p>
175      Let's recap how LLMs work, starting with their architecture and then moving onto inference mechanics. This will provide a starting point for an optimized implementation and help us establish benchmarks.
176    </p>
177    <p>
178      Almost<dt-fn>There are also state-space models like <a href="https://arxiv.org/abs/2312.00752">Mamba</a> which purport to be more efficient and scalable to long sequences than transformers, but they don't appear to have found much success outside of
179      niches like <a href="https://www.cartesia.ai/">low-power ML and non-discrete data domains like audio/video</a>.</dt-fn> every major open-weights LLM uses the same architecture<dt-fn>To the point where new foundation models explicitly state that they 
180      use a "standard architecture", with the value added being the training. This has led to some misunderstandings, see for instance <a href="https://blog.eleuther.ai/nyt-yi-34b-response/#how-all-llms-are-similar">https://blog.eleuther.ai/nyt-yi-34b-response/#how-all-llms-are-similar</a></dt-fn> 
181      (sequential transformer blocks), with some minor variations/innovations since GPT-2:
182      </p><ul class="compact-li">
183        <li><a href="https://arxiv.org/abs/2305.13245v3">Grouped query attention (and multi-query attention)</a></li>
184        <li><a href="https://arxiv.org/abs/2401.04088">Mixture-of-experts-based feedforward networks</a></li>
185        <li><a href="https://arxiv.org/abs/2002.05202">GLU-based instead of MLP-based feedforward networks</a></li>
186        <li><a href="https://arxiv.org/abs/2002.05202">Different activation functions for feedforward networks</a></li>
187        <li><a href="https://arxiv.org/abs/1910.07467">Different layer normalizations</a></li>
188        <li><a href="https://arxiv.org/abs/2104.09864">Rotary position embeddings</a></li>
189      </ul>
190    <p></p>
191    <p>
192      Loading models from different architectures is thus essentially defining a customizable transformer block class, then creating a sequence of these configured with the 
193      right bells and whistles and initializing them with the <a href="https://huggingface.co/docs/safetensors/en/index">safetensors</a> weights. This article will focus on just 
194      one architecture - Mistral v0.2 - but if you're curious you can read 
195      <a href="https://github.com/ggerganov/llama.cpp/blob/master/docs/development/HOWTO-add-model.md">how llama.cpp adds support for new models</a>.
196    </p>
197    <h2 id="section-1.1">1.1 Inference overview</h2>
198    <p>
199      At a high level, inference looks like the C++ pseudocode below:
200    </p>
201    <!-- <dt-code block language="clike"> -->
202      <pre class="l-middle">        <code class="language-clike">
203/* PSUEDOCODE */
204
205void generate(Model&amp; model, std::string prompt, int steps) {
206  std::vector&lt;int&gt; encoded = tokenizer.encode(prompt);
207  InferenceState s(model);
208  
209  // 1. Prefill step: Forward the model on each prompt token, discarding 
210  // the output. This lets the model read the prompt and hydrates the KV 
211  // cache.
212  for (int token : encoded) {
213    model.forward(s, token);
214  }
215  // 2. Decode step: Forward the model repeatedly, generating 1 token at a time.
216  for (int i = 0; i &lt; steps; i++) {
217    model.forward(s, encoded.back());
218    int next_token = sampler.sample(s.logits);
219    encoded.push_back(next_token);
220    std::cout &lt;&lt; tokenizer.decode_one(next_token) &lt;&lt; std::flush;
221    if (next_token == tokenizer.EOS) {
222      break;
223    }
224  }
225}
226        </code>
227    </pre>
228  <!-- </dt-code> -->
229    <p>
230      We can start to see differences between training and inference immediately. Inference - at least the local kind that we care about - is usually single batch. 
231      For prompt completion and use cases like generating essays, the “decode phase” takes up the majority of execution and involves computing attention between the past context 
232      and just a single token (or query timestep). 
233    </p>
234    <p>
235      The prefill step is more similar to training in that we're given a complete sequence to attend over, but more on that later. In chatbots there's also an “append” step when 
236      passing the model additional user messages which is like prefill, but I won't talk about that in this article as our implementation will support only completions. 
237    </p>
238    <p>
239      The model forward pass looks like so:
240    </p>
241    <!-- <dt-code block language="clike"> -->
242      <pre class="l-middle">        <code class="language-clike">
243/* PSUEDOCODE */
244
245// InferenceState is the minimum set of buffers needed to
246// hold state during the forward pass and exists to avoid
247// extra allocations
248void Model::forward(InferenceState&amp; s, int token) {
249  // The embedding table maps token IDs to embedding vectors,
250  // which are copied into a buffer of the inference state
251  s.x = copy_embedding(token, this-&gt;token_embedding_table);
252  
253  // Models consist of a sequence of transformer blocks which
254  // mutate the inference state in order
255  for (Block&amp; b : this-&gt;blocks) {
256    b-&gt;block(s);
257  }
258  
259  // Usually there is a layer norm right before the final classifier
260  s.x = layernorm(s.x, this-&gt;lm_head_prenorm_weights);
261  // Typically we end with a linear transform from (dim) -&gt; (vocab_size)
262  s.logits = linear(s.x, this-&gt;lm_head_classifier_weights); 
263}
264
265void Block::block(InferenceState&amp; s) {
266  s.x_resid = layernorm(s.x, this-&gt;att_prenorm_weights);
267  // Multi-head attention typically includes: 
268  // 1. RoPE on input (element-wise mutation w/ sines/cosines)
269  // 2. QKV matmuls and updating the KV cache
270  // 3. Causal self-attention, softmax, and value mixing
271  // 4. Projection back into the residual stream
272  s.x_resid = multi_head_attn(
273    s.x_resid,
274    this-&gt;wq, 
275    this-&gt;wk, 
276    this-&gt;wv, 
277    this-&gt;key_cache,
278    this-&gt;value_cache
279  );
280  s.x += s.x_resid;
281  s.x_resid = layernorm(s.x, this-&gt;ffn_prenorm_weights);
282  // On modern architectures like Llama, this is a GLU feedforward 
283  // with 3 linear transforms, not a simple MLP:
284  // -&gt; w2(F.silu(w1(x)) * w3(x))
285  // Some architectures also split the FFN into a mixture of experts.
286  s.x_resid = ffn(s.x_resid, this-&gt;w1, this-&gt;w2, this-&gt;w3);
287  s.x += s.x_resid;
288}
289        </code>
290    </pre>
291  <!-- </dt-code> -->
292    <p>
293      This should look roughly familiar, if expressed a bit more low-level in typical C++ fashion. 
294    </p>
295    <p>
296      The main thing worth noting is that unlike training, inference can use a KV cache to store past keys and values for each block. 
297      We keep things simple and implement this as a simple ring buffer (known as sliding window attention in the literature), which is sufficient to 
298      support exact attention up to some maximum context length. Some exact attention implementations like PagedAttention use more complex KV caches to 
299      improve aspects like memory footprint.
300    </p>
301    <h2 id="section-1.2">1.2 Bottlenecks and benchmarks</h2>
302    <p>
303      Now we are ready to discuss bottlenecks and benchmarks. First, a fact: inference is memory-bandwidth-bound on modern hardware. For more, see 
304      this <a href="https://zeux.io/2024/03/15/llm-inference-sol">excellent blog post</a> by Arseny Kapoulkine, but the gist of it is:
305      </p><ul class="compact-li">
306        <li>
307          Every time we generate a token we need to read the entire model, performing only a few floating point operations per weight.
308        </li>
309        <li>
310          Modern CPUs and GPUs are <i>extremely</i> fast at floating point operations. The key metric is the FLOPs/s-to-memory-bandwidth-ratio (FLOPs/byte). 
311          For instance, the AMD Ryzen 7950X has about a 40:1 ratio, while the RTX 4090 has an 82:1 ratio. The AMD EPYC 7702P on my server has a less impressive, 
312          but still significant 10:1 ratio.
313        </li>
314      </ul>
315    <p></p>
316    <p>
317      This is why <a href="https://huggingface.co/docs/optimum/en/concept_guides/quantization">model quantization</a> is so effective at improving inference speed. 
318      It's not just allowing the hardware to use faster instructions (which is sometimes true), 
319      but also shrinking the input that we need to fit through the bandwidth bottleneck.
320    </p>
321    <p>
322      We can use bandwidth to establish a theoretical “speed of light”, or the max token throughput we can achieve. On my machine with an AMD EPYC 7702P and RTX 4090:
323      </p><ul class="compact-li">
324        <li>
325          EPYC 7702P max bandwidth<dt-fn><a href="https://www.amd.com/content/dam/amd/en/documents/products/epyc/amd-epyc-7002-series-datasheet.pdf">See the official datasheet</a>.</dt-fn>: <code>204.8 GB/s</code>
326        </li>
327        <li>
328          RTX 4090 max bandwidth<dt-fn>See <a href="https://en.wikipedia.org/wiki/GeForce_40_series#Desktop">wikipedia</a>.</dt-fn>: <code>1008 GB/s</code>
329        </li>
330        <li>
331          <p>
332            Mistral-7B-Instruct-v0.2-FP32 with a 4k context window and FP32 KV-cache is <code>29516398592 bytes</code>
333          </p>
334          <ul class="compact-li">
335            <li>
336              <code>204.8e9 bytes/s / 29516398592 bytes/tok = ~6.9 tok/s</code> for EPYC 7702P
337            </li>
338            <li>
339              This won't fit in the 24GB of RTX 4090 VRAM so we'll skip it.
340            </li>
341          </ul>
342        </li>
343        <li>
344          <p>
345            Mistral-7B-Instruct-v0.2-FP16 with a 4k context window and FP16 KV-cache is <code>15020875776 bytes</code>
346          </p>
347          <ul class="compact-li">
348            <li>
349              <code>204.8e9 bytes/s / 15020875776 bytes/tok = ~13.6 tok/s</code> for EPYC 7702P
350            </li>
351            <li>
352              <code>1008e9 bytes/s / 15020875776 bytes/tok = ~67.1 tok/s</code> for RTX 4090
353            </li>
354          </ul>
355        </li>
356      </ul>
357    <p></p>
358    <p>
359      Note that how close we can actually come to the theoretical bounds varies a bit depending on the hardware. We fortunately have a few popular inference 
360      engines that we can look at to set more realistic targets. On my machine<dt-fn>Unfortunately I wasn't able to test calm CPU as my machine doesn't support the extensions 
361      needed for it to compile.</dt-fn> using Mistral-7B-Instruct-v0.2 in FP16 with 4k context I'm able to get:
362    </p>
363    <table>
364      <tbody><tr>
365        <th>Program</th>
366        <th>Avg. throughput (~120 tokens)</th>
367        <th>Avg. throughput (~4800 tokens)</th>
368      </tr>
369      <tr>
370        <td>llama.cpp, CPU<dt-fn>With number of threads tuned.</dt-fn></td>
371        <td>8.7 tok/s</td>
372        <td>7.6 tok/s</td>
373      </tr>
374      <tr>
375        <td>huggingface transformers, GPU<dt-fn>See <a href="#appendix">appendix</a> for benchmark code. For some reason, this was the highest variance of all. Shown are a best of 3 run.</dt-fn></td>
376        <td>25.9 tok/s</td>
377        <td>25.7 tok/s</td>
378      </tr>
379      <tr>
380        <td>llama.cpp, GPU</td>
381        <td>61.0 tok/s</td>
382        <td>58.8 tok/s</td>
383      </tr>
384      <tr>
385        <td>calm, GPU</td>
386        <td>66.0 tok/s</td>
387        <td>65.7 tok/s</td>
388      </tr>
389    </tbody></table>
390    <dt-byline></dt-byline>
391
392    <h1 id="section-2">2. Inference on the CPU</h1>
393    <p>
394      We begin with a naive implementation on CPU (the code is available <a href="https://github.com/andrewkchan/yalm/blob/ec8c8fec911794c788c50dfe5d42ae9e1ef0e905/src/infer.cpp#L240">here</a>). 
395      It's a straightforward single-threaded implementation with a 4k KV cache that only supports FP32 weights and no explicit SIMD of any kind. It achieves a blazing fast 
396      throughput of 0.6 tok/s. Here's what that looks like:
397    </p>
398    <div class="l-middle">
399      <video controls="" src="./yalm-assets/cpu-naive.mov" style="width: min(90vw, 720px)"></video>
400    </div>
401    <h2 id="section-2.1">2.1 Multithreading</h2>
402    <p>
403      The first optimization step we can do is to begin parallelizing our code on the thread level. Equipped with our handy OpenMP pragmas, we go hunting for embarrassingly parallel 
404      opportunities. We'll optimize the same spots as llama2.c, and I'll go over each one to show the improvement.
405    </p>
406    <p>
407      First, <a href="https://github.com/andrewkchan/yalm/commit/2a65dcbdff106976aeff1f08a037c6b6ece5b80b">adding a single line of code</a> parallelizes our widely matrix-vector multiplication 
408      function so that each thread handles a row of the output:
409    </p>
410    <!-- <dt-code block language="clike"> -->
411      <pre class="l-middle">        <code class="language-clike">
412static void matmul(float* xout, float* x, float* w, int n, int d) {
413  // W (d,n) @ x (n,) -&gt; xout (d,)
414  int i;
415#pragma omp parallel for private(i)
416  for (i = 0; i &lt; d; i++) {
417    float val = 0.0f;
418    for (int j = 0; j &lt; n; j++) {
419      val += w[i * n + j] * x[j];
420    }
421    xout[i] = val;
422  }
423}
424        </code>
425    </pre>
426  <!-- </dt-code> -->
427    <p>
428      This is a big improvement that takes us to 4.2 tok/s with a bit of tuning to find the right number of threads:
429    </p>
430    <div class="l-middle">
431      <video controls="" src="./yalm-assets/cpu-openmp-mm.mov" style="width: min(90vw, 720px)"></video>
432    </div>
433    <p>
434      Next, we can parallelize our multi-head attention computation (<a href="https://github.com/andrewkchan/yalm/compare/2a65dcbdff106976aeff1f08a037c6b6ece5b80b..f39081de410ced856a4ba301234ab3fc6948e4a4">code here</a>) 
435      so that each thread gets an attention head to compute. This is a less immediately clear improvement, but gets us to 
436      4.4 tok/s for short context generations, and likely much better for long contexts.
437    </p>
438    <div class="l-middle">
439      <video controls="" src="./yalm-assets/cpu-openmp-mha.mov" style="width: min(90vw, 720px)"></video>
440    </div>
441    <h2 id="section-2.2">2.2 Weight quantization and SIMD</h2>
442    <p>
443      The next potential opportunity is to use SIMD. The EPYC 7702P CPU supports AVX and AVX2, which let us work with 256-bit vectors of 8 packed float32 values at a time. 
444      In our ubiquitous <code>matmul</code> function, we could try loading, multiplying, and accumulating 8 values at a time in the inner loop, which would let each thread finish 
445      its row-column dot product up to 8 times faster!
446    </p>
447    <p>
448      Unfortunately, inspecting our compiled code via <code>objdump</code> reveals that <code>matmul</code> is in fact already using AVX instructions (notice <code>vmovups</code>) to perform a vectorized dot 
449      product in the case that the inputs are large enough. It seems GCC is too smart:
450    </p>
451    <!-- <dt-code block language="plaintext"> -->
452      <pre class="l-middle">        <code class="language-plaintext">
4531f5:       c4 e3 7d 19 c1 01       vextractf128 xmm1,ymm0,0x1
4541fb:       c5 f0 58 c0             vaddps xmm0,xmm1,xmm0
4551ff:       c5 f8 12 c8             vmovhlps xmm1,xmm0,xmm0
456203:       c5 f0 58 c8             vaddps xmm1,xmm1,xmm0
457207:       c5 f0 c6 c1 55          vshufps xmm0,xmm1,xmm1,0x55
45820c:       c5 f8 58 c1             vaddps xmm0,xmm0,xmm1
459        </code>
460    </pre>
461  <!-- </dt-code> -->
462    <p>
463      Let's turn to quantization. We won't explore the full gamut of quantized formats, since the goal of this article is to explore the breadth of optimizations, and we've 
464      chosen our benchmarks with fixed formats. Instead, we'll just quantize our weights to FP16, which is the bare minimum needed to get it loaded onto the RTX 4090 VRAM anyway.
465    </p>
466    <p>
467      One hiccup is that many CPUs do not support native float16 math. But barring that, we'd like to keep our calculations in float32 as much as possible anyway to mitigate 
468      effects on accuracy, and we should be able to without trading off performance as long as bandwidth remains a bottleneck. 
469    </p>
470    <p>
471      So instead we leverage the fact that many CPUs do still support converting float16 values to float32 via the F16C x86 extension (which has been well-supported for over a 
472      decade now) to load float16 weights and convert them to float32 just-in-time for calculations. Among other things, this requires us to explicitly vectorize the loads in 
473      our matmul function from before because GCC doesn't know how to handle the half-precision arrays:
474    </p>
475    <!-- <dt-code block language="clike"> -->
476      <pre class="l-middle">        <code class="language-clike">
477// F16C code technically operates on 16-bit unsigned short integers
478typedef uint16_t f16_t;
479
480// matmul supporting float16 weights via the F16C extension, which allows
481// conversion into float32 values before calculations.
482static void matmul(float* xout, float* x, f16_t* w, int n, int d) {
483#if defined(__AVX2__) &amp;&amp; defined(__F16C__)
484  // W (d,n) @ x (n,) -&gt; xout (d,)
485  assert(n % 16 == 0);
486  int i;
487#pragma omp parallel for private(i)
488  for (i = 0; i &lt; d; i++) {
489    // Vectorized dot product of w[i][:] and x[:] where w is a packed float16 array.
490    __m256 sumlo = _mm256_setzero_ps();
491    __m256 sumhi = _mm256_setzero_ps();
492    for (int j = 0; j &lt; n; j+=16) {
493      // Extract the next set of 16 float16 weights from `w` and store them
494      // to two separate float32 vectors of width 8 (`wveclo_ps`, `wvechi_ps`)
495      __m256i wvec = _mm256_loadu_si256((__m256i*)&amp;w[i * n + j]);
496      __m128i wveclo = _mm256_extractf128_si256(wvec, 0);
497      __m128i wvechi = _mm256_extractf128_si256(wvec, 1);
498      __m256 wveclo_ps = _mm256_cvtph_ps(wveclo);
499      __m256 wvechi_ps = _mm256_cvtph_ps(wvechi);
500      // Extract the next two float32 vectors of width 8 `xveclo`, `xvechi` from `x`
501      __m256 xveclo = _mm256_loadu_ps(&amp;x[j]);
502      __m256 xvechi = _mm256_loadu_ps(&amp;x[j + 8]);
503      // Compute vectorized FMAs: sumlo += wveclo * xveclo, sumhi += wvechi * xvechi
504      sumlo = _mm256_fmadd_ps(wveclo_ps, xveclo, sumlo);
505      sumhi = _mm256_fmadd_ps(wvechi_ps, xvechi, sumhi);
506    }
507    // Horizontally reduce width-8 float32 vectors sumlo, sumhi to a scalar.
508    __m256 sum8 = _mm256_add_ps(sumlo, sumhi);              // sum8[0:8] = sumlo[0:8] + sumhi[0:8]
509    __m128 sum4 = _mm_add_ps(                               // sum4[0:4] = sum8[0:4] + sum8[4:8]
510      _mm256_extractf128_ps(sum8, 0), 
511      _mm256_extractf128_ps(sum8, 1)
512    );
513    __m128 sum1 = _mm_dp_ps(sum4, _mm_set1_ps(1.0f), 0xf1); // sum1[0] = dot(sum4, [1,1,1,1])
514    xout[i] = _mm_cvtss_f32(sum1);
515  }
516#else
517  assert(false &amp;&amp; "float16 not supported on this platform");
518#endif
519}
520        </code>
521    </pre>
522  <!-- </dt-code> -->
523    <p>
524      <a href="https://github.com/andrewkchan/yalm/blob/462b3705a2dc75d2ae22e3b6a4c45d9b92bc20ea/src/infer.cpp">The resulting implementation</a> does not yield any difference in perplexity for short texts. It's also nearly twice as fast at 8.2-8.4 tok/s:
525    </p>
526    <div class="l-middle">
527      <video controls="" src="./yalm-assets/cpu-f16c.mov" style="width: min(90vw, 720px)"></video>
528    </div>
529
530    <dt-byline></dt-byline>
531
532    <h1 id="section-3">3. Inference on the GPU</h1>
533    <br>
534    <p>
535      Having quantized our model to half its size, we can now load it onto our RTX 4090 and begin a GPU inference implementation. 
536      Remember our rules: raw C++/CUDA only, no CUTLASS, cuBLAS, cuDNN, or other libraries.
537    </p>
538    <p>
539      If you haven't seen CUDA code before, see <a href="https://developer.nvidia.com/blog/easy-introduction-cuda-c-and-c/">An Easy Introduction to CUDA C and C++</a> for a good intro. 
540      At a high level, CUDA allows you to execute a C++ function (”kernel”) on the GPU in parallel over a <i>grid</i> of <i>threads</i> where:
541      </p><ul class="compact-li">
542        <li>
543          Each thread receives the same function arguments as all the others, but can make its own local variables, and is assigned its own <code>threadIdx</code> which can be used to determine 
544          what work it is responsible for.
545        </li>
546        <li>
547          Threads are additionally organized into blocks which have their own <code>blockIdx</code> and a fixed number of threads (<code>blockDim</code>). Threads in the same block can efficiently cooperate 
548          by sharing data through shared memory, which is faster than the global memory that all threads in a grid may otherwise access.
549        </li>
550        <li>
551          The number of blocks and threads per block is specified when invoking a kernel in triple-chevrons <code>&lt;&lt;&lt; numBlocks, threadsPerBlock &gt;&gt;&gt;</code> and may be given as <code>int</code> or <code>dim3</code>.
552        </li>
553      </ul>
554    <p></p>
555    <h2 id="section-3.1">3.1 A naive port to CUDA</h2>
556    <p>
557      <a href="https://github.com/andrewkchan/yalm/blob/ec8c8fec911794c788c50dfe5d42ae9e1ef0e905/src/infer.cpp">A first, naive implementation in 270 lines of C++/CUDA</a> attempts to translate our CPU operations 1-1 to kernels, with some extra kernels for things like vector additions. 
558      So we end up with a fairly long list of kernels, but more importantly, host C++ code in our GPU backend which looks almost like our CPU backend, 
559      but with <code>&lt;&lt;&lt; . &gt;&gt;&gt;</code> attached to the function calls. For instance, here's the half of our block forward pass 
560      in our GPU backend where we pass the activations through a feedforward network before adding back as residuals:
561    </p>
562    <!-- <dt-code block language="clike"> -->
563      <pre class="l-middle">        <code class="language-clike">
564// mix self.w2(F.silu(self.w1(x)) * self.w3(x))
565// Note this is a feedforward with a GLU, not a simple MLP.
566matmul&lt;&lt;&lt;c.hidden_dim, WARP_SIZE&gt;&gt;&gt;(
567  w1<float>(), s.xb(), c.dim, c.hidden_dim, s.hb()
568);
569matmul&lt;&lt;&lt;c.hidden_dim, WARP_SIZE&gt;&gt;&gt;(
570  w3<float>(), s.xb(), c.dim, c.hidden_dim, s.hb2()
571);
572glu_gelu&lt;&lt;&lt;
573  (c.hidden_dim + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK, 
574  MAX_THREADS_PER_BLOCK,
575&gt;&gt;&gt;(
576  s.hb(), s.hb2(), s.hb()
577);
578matmul&lt;&lt;&lt;c.dim, WARP_SIZE&gt;&gt;&gt;(
579  w2<float>(), s.hb(), c.hidden_dim, c.dim, s.xb2()
580);
581// ffn residual back into x
582add_residuals&lt;&lt;&lt;
583  (c.dim + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK,
584  MAX_THREADS_PER_BLOCK
585&gt;&gt;&gt;(
586  s.x(), s.xb2(), c.dim, s.x()
587);
588        </float></float></float></code>
589    </pre>
590  <!-- </dt-code> -->
591    <p>
592      And the corresponding CPU code:
593    </p>
594    <!-- <dt-code block language="clike"> -->
595      <pre class="l-middle">        <code class="language-clike">
596// mix self.w2(F.silu(self.w1(x)) * self.w3(x))
597// Note this is a feedforward with a GLU, not a simple MLP.
598matmul(s.hb(), s.xb(), w1&lt;T&gt;(), c.dim, c.hidden_dim);
599matmul(s.hb2(), s.xb(), w3&lt;T&gt;(), c.dim, c.hidden_dim);
600for (int i = 0; i &lt; c.hidden_dim; ++i) {
601  s.hb()[i] = gelu(s.hb()[i]) * s.hb2()[i];
602}
603matmul(s.xb2(), s.hb(), w2&lt;T&gt;(), c.hidden_dim, c.dim);
604// residual connection back into x
605for (int i = 0; i &lt; c.dim; ++i) {
606  s.x()[i] += s.xb2()[i];
607}
608        </code>
609    </pre>
610  <!-- </dt-code> -->
611    <p>
612      This works because even though CUDA kernels are executed asynchronously, kernels on the same stream (or the default stream) will never execute 
613      concurrently; one starts only when the previous finishes all threads. So we just need to have our final kernel write to device-mapped memory 
614      on the host and perform a deviceSync at the end of the forward pass so that the output is available for sampling. 
615    </p>
616    <h2 id="section-3.2">3.2 Better matmuls</h2>
617    <p>
618      An interesting deep dive here is our <code>matmul</code> kernel. We saw earlier that this function was a huge piece of runtime on the CPU, and optimizing it 
619      via OpenMP yielded huge wins. So let's make sure <code>matmul</code> is well-optimized before proceeding.
620    </p>
621    <p>
622      A naive implementation of <code>matmul</code> might look something like our OpenMP-parallelized CPU code, where we have each GPU thread handle computing 1 row (element) of the resulting vector:
623    </p>
624    <!-- <dt-code block language="clike"> -->
625      <pre class="l-middle">        <code class="language-clike">
626__global__
627void matmul(const float* A, const float* x, int n, int d, float* out) {
628  // A (d,n) @ x (n,) -&gt; out (d,)
629  int i = blockIdx.x * blockDim.x + threadIdx.x;
630  if (i &gt;= d) return;
631  float sum = 0.0;
632  for (int j = 0; j &lt; n; j++) {
633    sum += A[n * i + j] * x[j];
634  }
635  out[i] = sum;
636}
637
638/* usage */
639int MAX_THREADS_PER_BLOCK = 1024;
640matmul&lt;&lt;&lt;
641  (d + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK, 
642  MAX_THREADS_PER_BLOCK
643&gt;&gt;&gt;(A, x, n, d, out);
644        </code>
645    </pre>
646  <!-- </dt-code> -->
647    <img src="./yalm-assets/mm-naive.png" class="l-middle" style="width: min(90vw, 650px)">
648    <p>
649      One big problem with this approach is that it will under-utilize our CUDA cores. Mistral-7B has a transformer input/output dimension of 4096, so if we're 
650      computing (for example) the last matmul before the output, we'll spin up 4096 threads. But an RTX 4090 can have 16384 simultaneous threads! Many of our cores 
651      will be sitting idle and we won't reach full FLOPs/s.
652    </p>
653    <p>
654      In addition to FLOPs, there's also memory load coalescing issues with this kernel - but more on that later. Suffice to say that substituting this kernel in 
655      for the one that we end up with leads to a throughput of 2.9 tok/s - slower than our CPU backend!
656    </p>
657    <p>
658      A better idea is to have one block per row. Threads within a block can efficiently cooperate, so we'll be able to leverage more threads to speed up computing the 
659      result for a single row. In this setup, each block will have exactly 1 w
659arp (a smaller group of threads within a block which are a fundamental execution unit on the hardware, 
660      with their own <a href="https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/">cooperation primitives</a>), then we'll use a warp-stride loop to 
661      sum across the entire row. At the end, we can perform a <a href="https://developer.nvidia.com/blog/faster-parallel-reductions-kepler/">warp sum reduction</a> 
662      to combine results from all threads:
663    </p>
664    <!-- <dt-code block language="clike"> -->
665      <pre class="l-middle">        <code class="language-clike">
666__device__ 
667inline float warp_reduce_sum(float val) {
668  for (int offset = WARP_SIZE / 2; offset &gt; 0; offset /= 2)
669    val += __shfl_down_sync(0xffffffff, val, offset);
670
671  return val;
672}
673
674__device__
675inline float matmul_row(const float* row, const float* x, int offset, int dim) {
676  float sum = 0.0;
677  for (int j = offset; j &lt; dim; j += WARP_SIZE) {
678    float v = row[j] * x[j];
679    sum += v;
680  }
681  return warp_reduce_sum(sum);
682}
683
684__global__
685void matmul(const float* A, const float* x, int n, int d, float* out) {
686  // A (d,n) @ x (n,) -&gt; out (d,)
687  // PRECOND: Blocks are 1-D and same size as warp.
688  int i = blockIdx.x;
689  if (i &gt;= d) return;
690  int offset = threadIdx.x;
691  float rowSum = matmul_row(&amp;A[n * i], x, offset, n);
692  if (threadIdx.x == 0) {
693    out[i] = rowSum;
694  }
695}
696
697/* usage */
698int BLOCK_SIZE = WARP_SIZE;
699matmul&lt;&lt;&lt;d, BLOCK_SIZE&gt;&gt;&gt;(A, x, n, d, out);
700        </code>
701    </pre>
702  <!-- </dt-code> -->
703    <img src="./yalm-assets/mm-better.png" class="l-middle" style="width: min(90vw, 650px)">
704    <p>
705      We can adapt the above kernel to blocks with more than one warp (it's easiest to keep 1 warp per row) as well. I will leave this and the question 
706      of why that might be a good idea as an exercise for the reader. :)
707    </p>
708    <p>
709      This kernel has much better thread utilization and memory read coalescing. With it configured to 1 warp per block, even with the near-1-1 arrangement 
710      of functions to kernels, we get a throughput of 51.7 tok/s - pretty good performance for a preliminary implementation! That said, we still have a ways 
711      to go before catching up to llama.cpp or calm - what can we improve?
712    </p>
713    <h2 id="section-3.3">3.3 Fusing and even better matmuls</h2>
714    <p>
715      Profiling a generation with <code>nsys</code> shows a few interesting things. First, we see that our GPU is being used for nearly the entire generation time, despite 
716      the host-device syncs at the end of every forward pass. There are some gaps in the device thread, indicating occasional idle time, but they are on the 
717      order of microseconds, which means we aren't ever CPU-bound:
718    </p>
719    <img src="./yalm-assets/nsys-gaps.png" class="l-middle" style="width: min(90vw, 650px)">
720    <p>
721      Second, we can see that 94.4% of our kernel executions are matrix multiplications:
722    </p>
723    <img src="./yalm-assets/nsys-mostly-matmul.png" class="l-middle" style="width: min(90vw, 650px)">
724    <p>
725      This means that if we were to get rid of every other kernel, we'd at best run in 94.4% of our current time, which would take us from 51.7 tok/s to 54.8 tok/s 
726      throughput, which is not quite at our goal. So we still need to optimize our matrix multiplications.
727    </p>
728    <p>
729      That said, we can approximate getting rid of other kernels by fusing them together. In particular, there are a few kernels that we can fuse into the nearest matmul, 
730      and a few matmuls that we can fuse together. For example, we can fuse together <code>matmul</code> and <code>add_residuals</code> into a single <code>fused_matmul_add_residuals</code> 
731      that directly sums to the destination. The Mistral architecture also performs a gated sum of 2 matmuls at some point: <code>F.silu(w1(x)) * w3(x)</code> - and we can do 
732      all of this in one kernel as all operations have 1 thread write to 1 element of the output vector, and all have the same dimension output.
733    </p>
734    <p>
735      This can let us “blend” together 2 dependent operations. For the first example, <code>add_residuals</code> could not begin until all threads of <code>matmul</code> finished, 
736      which is problematic if individual threads of <code>matmul</code> take longer. Additionally we avoid one extra read and write from each thread to global memory. 
737    </p>
738    <p>
739      For the curious, the full change is <a href="https://github.com/andrewkchan/yalm/commit/7e3bb32a91117c9892363e8d2536b90cc2b1b2f3">here</a>. This takes us to 54.1 tok/s!
740    </p>
741    <p>
742      Now we can return to optimizing <code>matmul</code>. From profiling via <code>ncu</code>, <code>matmul</code> writes are not optimally coalesced. The relevant warning diagnostic 
743      is shown below:
744    </p>
745    <!-- <dt-code block language="plaintext"> -->
746      <pre class="l-middle">        <code class="language-plaintext">
747    Section: Memory Workload Analysis
748    --------------------------- ------------ ------------
749    Metric Name                  Metric Unit Metric Value
750    --------------------------- ------------ ------------
751    Memory Throughput           Gbyte/second       533.22
752    Mem Busy                               %        24.82
753    Max Bandwidth                          %        90.26
754    L1/TEX Hit Rate                        %        65.94
755    L2 Compression Success Rate            %            0
756    L2 Compression Ratio                                0
757    L2 Hit Rate                            %         2.03
758    Mem Pipes Busy                         %        28.33
759    --------------------------- ------------ ------------
760
761    Section: Memory Workload Analysis Tables
762    ...
763    ----- --------------------------------------------------------------------------------------------------------------
764    WRN   The memory access pattern for stores from L1TEX to L2 is not optimal. The granularity of an L1T
764EX request to  
765          L2 is a 128 byte cache line. That is 4 consecutive 32-byte sectors per L2 request. However, this kernel only  
766          accesses an average of 1.0 sectors out of the possible 4 sectors per cache line. Check the Source Counters    
767          section for uncoalesced stores and try to minimize how many cache lines need to be accessed per memory        
768          request.                                                                                                      
769    ----- --------------------------------------------------------------------------------------------------------------
770    ...
771        </code>
772    </pre>
773  <!-- </dt-code> -->
774    <p>
775      At this point it's helpful to define what it means for a load or store to global memory to be coalesced in CUDA. When 
776      threads in the same warp issue loads or stores to global memory, the loads or stores can be grouped together and 
777      performed as a single transaction if they are appropriately aligned and refer to consecutive regions of memory 
778      (see the diagrams in this <a href="https://stackoverflow.com/a/5044424/4151721">excellent SO post</a>).
779    </p>
780    <p>
781      This is important because GPU global memory transactions can only be done at certain levels of granularity (32, 64, or 
782      128 bytes). So if 32 threads in the same warp are reading or writing to distinct 4-byte floats, if the access is coalesced, 
783      a single transaction of 128 bytes will be performed, whereas if the access is not coalesced, 32 transactions of 32-bytes 
784      each may be performed, wasting a ton of bandwidth.
785    </p>
786    <p>
787      In our <code>matmul</code> kernel, reads are coalesced, but writes are not. We have 1 warp per block with each warp handling 1 output element, 
788      so we are issuing 1 write per block of minimum size 32 bytes but only updating a single 4-byte element. Increasing the number of 
789      warps per block does not help because coalescing is done at the warp level, not the block level. 
790    </p>
791    <p>
792      Instead, we can have each warp compute their own result, but then collect all results into lanes of the first warp, then have the first warp 
793      issue a single coalesced write. The block-wide collection can be done fast through shared memory. The kernel:
794    </p>
795    <!-- <dt-code block language="clike"> -->
796      <pre class="l-middle">        <code class="language-clike">
797__device__ inline float blocktranspose(float v, float def) {
798  // Performs block-and-warp transpose operation:
799  //   For a block containing K warps where lane 0 contains val_k,
800  //   this function returns:
801  //   - For warp 0, lane K: val_k
802  //   - For all other warps and lanes: def
803  int lane = threadIdx.x % warpSize;
804  int warp = threadIdx.x / warpSize;
805  
806  // Will hold results of all warps.
807  // Each lane of the warp accumulates across 1 head element at a time.
808  // NOTE: Assumes warpSize is 32
809  __shared__ float sm[32];
810  if (lane == 0) sm[warp] = v;
811  __syncthreads();
812  
813  return lane &lt; blockDim.x / warpSize ? sm[lane] : def;
814}
815
816template &lt;typename T&gt;
817__global__
818void matmul_wide(const T* A, const float* x, int n, int d, float* out) {
819  // A (d,n) @ x (n,) -&gt; out (d,)
820  // PRECOND: Block is 1-D and contains WPB warps.
821  int i = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;
822  if (i &gt;= d) return;
823  // Warp j computes sum for row at &lt;blockIdx.x*WPB + j&gt;
824  // Lane 0 of each warp will hold result
825  int k = threadIdx.x % warpSize;
826  float rowSum = matmul_row(&amp;A[n * i], x, k, n);
827  // Transpose values so lane k in warp 0 contains row at &lt;blockIdx.x*WPB + k&gt;
828  // For WPB=32, this allows us to coalesce 32 float32 writes into a single 128-byte store
829  rowSum = blocktranspose(rowSum, 1.0);
830  if (threadIdx.x &lt; blockDim.x / warpSize) {
831    int block_start_i = blockIdx.x * blockDim.x / warpSize;
832    out[block_start_i + k] = rowSum;
833  }
834}
835        </code>
836    </pre>
837  <!-- </dt-code> -->
838    <p>
839      Profiling this with <code>ncu</code> shows a nearly 10% improvement (from 600 usecs to 550 usecs) on toy matrix dimensions. 
840      After adapting the final LM classifier and block <code>fused_matmul_add_residual</code> kernels to use this improved coalescing, 
841      we reach 56.1 tok/s throughput:
842    </p>
843    <div class="l-middle">
844      <video controls="" src="./yalm-assets/gpu-fused-better-matmuls.mov" style="width: min(90vw, 720px)"></video>
845    </div>
846
847    <h2 id="section-3.4">3.4 Attention and long context generation</h2>
848    <p>
849      One major opportunity for improvement is in long context generations. When generating 5k tokens, our model's 
850      average throughput degrades to ~48 tok/s and attention kernels go from ~5% to &gt;10% of runtime. Compare this to llama.cpp and calm, 
851      which are able to maintain throughput in the high 50s or greater. What's going on?
852    </p>
853    <p>
854      With short context lengths, attention is fast - as there are not too many tokens to communicate over - and decoding is instead dominated by other, 
855      non-context sensitive operations such as generating the q/k/v vector representations of the input and passing the attention block output through the 
856      feedforward network.
857    </p>
858    <p>
859      The story is different with long contexts. This is constrained with sliding window attention, but as window size approaches hidden dimension, I'd 
860      expect attention runtime to approach that of our feedforward network.
861    </p>
862    <p>
863      More precisely, let's consider how 2 attention operations (computing the attention scores, and using them to mix values) compare to a single matrix 
864      multiply from dim to hidden_dim (the feedforward, which in modern architectures is based on a GLU rather than a simple MLP, should effectively contain 3 of these). 
865      We have...
866      </p><ul class="compact-li">
867        <li>
868          Matmul going from <code>dim</code> to <code>hidden_dim</code>: <code>n_heads</code> dot products of size <code>head_dim</code> per row, <code>hidden_dim</code> rows
869          <ul class="compact-li">
870            <li><code>n_heads * hidden_dim</code> dot products of size <code>head_dim</code></li>
871          </ul>
872        </li>
873        <li>
874          Attention over sequence of length <code>seq_len</code>:
875          <ul class="compact-li">
876            <li><code>n_heads * seq_len</code> dot products of size <code>head_dim</code></li>
877          </ul>
878        </li>
879        <li>
880          Attention mix over sequence of length <code>seq_len</code>:
881          <ul class="compact-li">
882            <li><code>n_heads * seq_len</code> multiply-adds (of a vector of size <code>head_dim</code>)</li>
883          </ul>
884        </li>
885      </ul>
886    <p></p>
887    <p>
888      In our benchmark, we use Mistral-v0.2 7B with <code>seq_len=4096</code> and <code>hidden_dim=14336</code>, so we'd expect our FFN matmuls to have roughly 3.5x the runtime of 
889      our attention kernels.
890    </p>
891    <p>
892      We can see that this is about right for <code>attn</code> (our scoring kernel) which takes about ~50 μs vs. ~150 μs for the final FFN matmul:
893    </p>
894    <img src="./yalm-assets/long-context-profile-attn.png" class="l-middle" style="width: min(90vw, 700px)">
895    <img src="./yalm-assets/long-context-profile-ffn.png" class="l-middle" style="width: min(90vw, 700px)">
896    <p>
897      However, something weird is going on with our <code>att_mix</code> kernel, which is taking ~150 μs, the same amount of time as our FFN matmuls!
898    </p>
899    <img src="./yalm-assets/long-context-profile-att-mix.png" class="l-middle" style="width: min(90vw, 700px)">
900    <p>
901      Profiling <code>att_mix</code> alone with <code>ncu</code> reveals that the kernel has serious memory bandwidth issues - it only uses 8% of memory throughput!
902    </p>
903    <!-- <dt-code block language="plaintext"> -->
904      <pre class="l-middle">        <code class="language-plaintext">
905  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
906    Section: GPU Speed Of Light Throughput
907    ----------------------- ------------- ------------
908    Metric Name               Metric Unit Metric Value
909    ----------------------- ------------- ------------
910    DRAM Frequency          cycle/nsecond         6.27
911    SM Frequency            cycle/nsecond         1.33
912    Elapsed Cycles                  cycle      5765045
913    Memory Throughput                   %         8.81
914    DRAM Throughput                     %         1.68
915    Duration                      msecond         4.35
916    L1/TEX Cache Throughput             %         5.17
917    L2 Cache Throughput                 %         8.81
918    SM Active Cycles                cycle   5685841.81
919    Compute (SM) Throughput             %         0.47
920    ----------------------- ------------- ------------
921
922    WRN   This kernel exhibits low compute throughput and memory bandwidth utilization relative to the peak performance 
923          of this device. Achieved compute throughput and/or memory bandwidth below 60.0% of peak typically indicate    
924          latency issues. Look at Scheduler Statistics and Warp State Statistics for potential reasons.                                                                                       
925      </code>
926    </pre>
927    <!-- </dt-code> -->
928    <p>
929      What's going on here? First let's review what this kernel is supposed to do, then the kernel code itself.
930      </p><ul class="compact-li">
931        <li>
932          Given 2 tensors in row-major order:
933          <ul class="compact-li">
934            <li>
935              <code>att</code> - attention scores with shape <code>(n_heads, kv_len)</code>
936            </li>
937            <li>
938              <code>vb</code> - value vectors with shape <code>(max_seq_len, n_kv_heads, head_dim)</code>
939            </li>
940          </ul>
941        </li>
942        <li>
943          We want to output a tensor <code>out</code> of shape <code>(n_heads, head_dim)</code>where <code>out[q] = att[q, :] @ vb[:, q//G, :]</code>where <code>G = n_heads//n_kv_heads</code> 
944          is the size of a group in grouped-query attention.
945        </li>
946        <li>
947          Below, we assume the number of timesteps in our KV cache <code>kv_len</code> is equal to <code>max_seq_len</code>:
948        </li>
949      </ul>
950    <p></p>
951    <img src="./yalm-assets/att_mix_math.png" class="l-middle" style="width: min(90vw, 700px)">
952    <p>
953      Below is the naive kernel that we've been using to do this. The idea is to spin up a block for every element <code>out[q, i]</code> that we need to compute. 
954      The threads of the block distribute over and mix the range of values <code>vb[:, q//G, i]</code> using a block-stride loop, then combine their results with a block sum 
955      reduction (below, we assume every block is the size of a warp so we can use efficient more warp reductions):
956    </p>
957    <!-- <dt-code block language="clike"> -->
958      <pre class="l-middle">        <code class="language-clike">
959__device__ 
960inline float warp_reduce_sum(float val) {
961  for (int offset = WARP_SIZE / 2; offset &gt; 0; offset /= 2)
962    val += __shfl_down_sync(0xffffffff, val, offset);
963
964  return val;
965}
966
967__global__
968void att_mix(
969  const float* vb, // (max_seq_len, n_kv_heads, head_dim) 
970  const float* att, // (n_heads, kv_len)
971  int head_dim, 
972  int n_heads, 
973  int n_kv_heads,
974  int seq_len, 
975  int max_seq_len, 
976  float* out // (n_heads, head_dim)
977) {
978  // PRECOND: blocks are 1-D and blockDim.x == WARP_SIZE
979  int h = blockIdx.x;
980  int group_size = n_heads / n_kv_heads;
981  int g = h / group_size;
982  int i = blockIdx.y;
983  int offset = threadIdx.x;
984  int kv_stride = n_kv_heads * head_dim;
985  
986  const float* atth = att + max_seq_len * h;
987  const float* vh = vb + head_dim * g;
988  float* outh = out + head_dim * h;
989  
990  float sum = 0.0;
991  for (int t = offset; t &lt; seq_len; t += WARP_SIZE) {
992    sum += vh[kv_stride * t + i] * atth[t];
993  }
994  sum = warp_reduce_sum(sum);
995  if (offset == 0) outh[i] = sum;
996}
997
998/* usage */
999dim3 tpb;
1000tpb.x = WARP_SIZE;
1001dim3 blocks;
1002blocks.x = n_heads;
1003blocks.y = head_dim;
1004att_mix&lt;&lt;&lt;blocks, tpb&gt;&gt;&gt;(
1005  vb, att, 
1006  head_dim, n_heads, n_kv_heads, 
1007  seq_len, max_seq_len, out
1008);
1009      </code>
1010    </pre>
1011    <!-- </dt-code> -->
1012    <p>
1013      Visualized below, we can see one big problem - the threads of each block are not accessing contiguous memory and so cannot coalesce any of their loads!
1014    </p>
1015    <img src="./yalm-assets/att_mix_naive.png" class="l-middle" style="width: min(90vw, 700px)">
1016    <p>
1017      Here's a first shot at improving the coalescing:
1018      </p><ul class="compact-li">
1019        <li>
1020          Have multiple blocks write to a single output element <code>out[h, i]</code>.
1021        </li>
1022        <li>
1023          Partition the sequence into contiguous time chunks for each block to handle. This can be done using the y-dimension of our grid. However, we still want the x-dimension to correspond to heads.
1024        </li>
1025        <li>
1026          Blocks should handle multiple output elements. 1 warp per block. <code>i</code> is determined by warp ID. This allows warps to coalesce loads.
1027        </li>
1028        <li>
1029          Threads use <code>atomicAdd</code> to accumulate result to <code>out[h, i]</code>.
1030        </li>
1031      </ul>
1032    <p></p>
1033    <p>
1034      In code:
1035    </p>
1036    <!-- <dt-code block language="clike"> -->
1037      <pre class="l-middle">        <code class="language-clike">
1038__global__
1039void att_mix(
1040  const float* vb,  // (max_seq_len, n_kv_heads, head_dim) 
1041  const float* att, // (n_heads, kv_len)
1042  int head_dim, 
1043  int n_heads, 
1044  int n_kv_heads,
1045  int seq_len, 
1046  int max_seq_len, 
1047  float* out // (n_heads, head_dim)
1048) {
1049  // PRECOND: blocks are 1-D and `out` has been zeroed
1050  int h = blockIdx.x;
1051  int group_size = n_heads / n_kv_heads;
1052  int g = h / group_size;
1053  int kv_stride = n_kv_heads * head_dim;
1054  
1055  const float* atth = att + max_seq_len * h;
1056  const float* vh = vb + head_dim * g;
1057  float* outh = out + head_dim * h;
1058  
1059  int t_per_thread = seq_len / gridDim.y;
1060  int t_start = blockIdx.y * t_per_thread;
1061  
1062  for (int i = threadIdx.x; i &lt; head_dim; i += blockDim.x) {
1063    float sum = 0.0;
1064    for (int t = t_start; t &lt; t_start + t_per_thread; t++) {
1065      sum += vh[kv_stride * t + i] * atth[t];	
1066    }
1067    atomicAdd(&amp;outh[i], sum);
1068  }
1069}
1070
1071int max_t_per_thread = 256;
1072
1073dim3 tpb;
1074tpb.x = warp_size;
1075dim3 blocks;
1076blocks.x = c.n_heads;
1077blocks.y = (kv_len + max_t_per_thread - 1) / max_t_per_thread;
1078cudaMemset(s.xb2(), 0, c.n_heads * c.head_dim * sizeof(float));
1079att_mix&lt;&lt;&lt;blocks, tpb&gt;&gt;&gt;(
1080  vb, s.att(),
1081  c.head_dim, c.n_heads, c.n_kv_heads, 
1082  kv_len, c.max_seq_len, s.xb2()
1083);
1084      </code>
1085    </pre>
1086    <!-- </dt-code> -->
1087    <p>
1088      Drawn out, we can see that the threads of each block are now loading contiguous elements of <code>vb</code> with each iteration of their inner loop:
1089    </p>
1090    <img src="./yalm-assets/att_mix_better.png" class="l-middle" style="width: min(90vw, 700px)">
1091    <p>
1092      Substituting this kernel speeds things up a lot for both long and (surprisingly) short context generations - we go from ~56 tok/s and ~48 tok/s respectively to 
1093      ~63 tok/s and ~57 tok/s! Profiles also reveal that <code>att_mix</code> now takes around the same time as <code>attn</code>, as we expect.
1094      Unfortunately, long context generations also noticeably degrade in both subjective quality and perplexity (arou
1094nd ~5x increase for one 4000-token text). 
1095    </p>
1096    <p>
1097      As it turns out, the problem is that <code>atomicAdd</code> operations to global memory sneakily don't add very small float values (on the order of <code>1e-40</code>) as 
1098      they're considered <a href="https://stackoverflow.com/questions/20189225/cuda-atomicaddfloat-does-not-add-very-small-values">subnormal values which global atomics always flush to zero</a>. 
1099      According to this <a href="https://forums.developer.nvidia.com/t/atomicadd-float-does-not-add-very-small-values/31531/8?u=akc">2013 forum post</a>, this shouldn't happen with 
1100      <code>atomicAdd</code> ops to shared memory, so here's a re-write of our kernel to accumulate results in block shared memory 
1101      and write to global memory when finished:
1102    </p>
1103    <!-- <dt-code block language="clike"> -->
1104      <pre class="l-middle">        <code class="language-clike">
1105__global__
1106void att_mix(
1107  const float* vb,  // (max_seq_len, n_kv_heads, head_dim) 
1108  const float* att, // (n_heads, kv_len)
1109  int head_dim, 
1110  int n_heads, 
1111  int n_kv_heads,
1112  int seq_len, 
1113  int max_seq_len, 
1114  float* out // (n_heads, head_dim)
1115) {
1116  // PRECOND: blocks are 2-D (warp_size, t_stride)
1117  int h = blockIdx.x;
1118  int group_size = n_heads / n_kv_heads;
1119  int g = h / group_size;
1120  int kv_stride = n_kv_heads * head_dim;
1121  
1122  const float* atth = att + max_seq_len * h;
1123  const float* vh = vb + head_dim * g;
1124  float* outh = out + head_dim * h;
1125  
1126  int warp_id = threadIdx.y;
1127  int t_stride = blockDim.y;
1128  
1129  // Capacity 32 since there can be at most 32 warps in a block.
1130  __shared__ float shared[32];
1131  
1132  for (int i = threadIdx.x; i &lt; head_dim; i += warpSize) {
1133    if (warp_id == 0) {
1134      shared[threadIdx.x] = 0;
1135    }
1136    __syncthreads();
1137    float sum = 0.0;
1138    for (int t = warp_id; t &lt; seq_len; t += t_stride) {
1139      sum += vh[kv_stride * t + i] * atth[t];	
1140    }
1141    atomicAdd(&amp;shared[threadIdx.x], sum);
1142    __syncthreads();
1143    if (warp_id == 0) {
1144      outh[i] = shared[threadIdx.x];
1145      shared[threadIdx.x] = 0;
1146    }
1147  }
1148}
1149
1150dim3 tpb;
1151tpb.x = warp_size;
1152tpb.y = min(kv_len, max_threads_per_block / warp_size);
1153dim3 blocks;
1154blocks.x = c.n_heads;
1155att_mix&lt;&lt;&lt;blocks, tpb&gt;&gt;&gt;(
1156  vb, s.att(),
1157  c.head_dim, c.n_heads, c.n_kv_heads, 
1158  kv_len, c.max_seq_len, s.xb2()
1159);
1160      </code>
1161    </pre>
1162    <!-- </dt-code> -->
1163    <p>
1164      This fixes our quality issues while maintaining the same speed at ~63.7 tok/s for a short generation:
1165    </p>
1166    <div class="l-middle">
1167      <video controls="" src="./yalm-assets/gpu-att-mix.mov" style="width: min(90vw, 720px)"></video>
1168    </div>
1169    <p>
1170      This finally gets us past llama.cpp's short generation throughput of 61.0 tok/s (see <a href="#section-1.2">§1.2</a>)!
1171    </p>
1172
1173    <h2 id="section-3.5">3.5 KV quantization and compiler gotchas</h2>
1174    <p>
1175      This last section is about KV cache quantization, which is a common optimization we haven't done yet, but also about some of 
1176      the optimizations that the <code>nvcc</code> CUDA compiler does for you automatically, what can happen when they break, and how to fix them. 
1177      Earlier we saw how automatic vectorization needed to be made explicit when working with FP16 weights in our CPU backend, and as 
1178      we'll see, switching to an FP16 KV cache will require similar work.
1179    </p>
1180    <p>
1181      First, a bit about KV cache quantization:
1182      </p><ul class="compact-li">
1183        <li>
1184          In our benchmark setting earlier (<a href="#section-1.2">§1.2</a>), llama.cpp and calm were actually using FP16 KV cache entries (because that is 
1185          their default setting), and we calculated the speed-of-light assuming the same. 
1186        </li>
1187        <li>
1188          Quantizing KV cache entries from FP32 to FP16 will not yield as large of a win as quantizing our weights did, since the KV cache 
1189          comprises a much smaller part of the total memory, but we still expect a win nonetheless as the total memory read each forward pass 
1190          will go from ~15.5 GB to ~15.0 GB. The effect should be largest for long contexts and our attention kernels, as the KV cache is 
1191          only used in attention.
1192        </li>
1193      </ul>
1194    <p></p>
1195    <p>
1196      Just switching our KV cache from FP32 to FP16 (while still performing computations in FP32, just like weight quantization) is simple 
1197      and boils down to replacing <code>float*</code> with <code>half*</code> and inserting the <code>__half2float</code> intrinsic in a 
1198      few key places: <a href="https://github.com/andrewkchan/yalm/commit/19a21912f2bd2e5ff0eb59bf5f3df1756038e3e0">see the code change here</a>. 
1199      Ensuring that it's performant is another story.
1200    </p>
1201    <p>
1202      Rather than make things faster, the naive patch above actually causes things to get slower. Long context generations are especially 
1203      impacted, with throughput down to ~53.6 tok/s from ~57 tok/s, and the <code>att_mix</code> kernel going from 5.5% (29 μs avg.) to a whopping 13.3% 
1204      of runtime (78 μs).
1205    </p>
1206    <img src="./yalm-assets/kernel-stats-kv32.png" class="l-middle" style="width: min(90vw, 700px)">
1207    <p class="caption">
1208      <i>nsys profile of a long context generation with FP32 KV cache</i>
1209    </p>
1210    <img src="./yalm-assets/kernel-stats-kv16.png" class="l-middle" style="width: min(90vw, 700px)">
1211    <p class="caption">
1212      <i>The same profile with our naive patch switching to FP16 KV cache</i>
1213    </p>
1214    <p>
1215      An <code>ncu</code> profile on a toy example reveals that memory throughput for the att_mix kernel has gotten much worse, 
1216      going from ~25.7% of the memory speed of light to ~6.8% and 140 μs to 309 μs. Why? One issue may be that we are simply issuing 
1217      smaller memory transactions than we previously did in each iteration of our loop. Recall that in <code>att_mix</code>, each thread accumulates 
1218      a sum over timesteps <code>t</code> like so:
1219    </p>
1220    <!-- <dt-code block language="clike"> -->
1221    <pre class="l-middle">      <code class="language-clike">
1222float sum = 0.0;
1223for (int t = warp_id; t &lt; seq_len; t += t_stride) {
1224  sum += vh[kv_stride * t + i] * atth[t];	
1225}
1226// atomicAdd `sum` when we're done
1227      </code>
1228    </pre>
1229    <!-- </dt-code> -->
1230    <p>
1231      <code>vh</code> is now a half-precision array, so loads from it must be converted to FP32 before being used. Our naive 
1232      patch does this by simply inserting a <code>__half2float</code> call:
1233    </p>
1234    <!-- <dt-code block language="clike"> -->
1235      <pre class="l-middle">        <code class="language-clike">
1236float sum = 0.0;
1237for (int t = warp_id; t &lt; seq_len; t += t_stride) {
1238  sum += __half2float(vh[kv_stride * t + i]) * atth[t];	
1239}
1240// atomicAdd `sum` when we're done
1241        </code>
1242      </pre>
1243    <!-- </dt-code> -->
1244    <p>
1245      This results in lower throughput because even though loads are still being coalesced in each warp for each loop iteration, 
1246      they are coalesced to 32 16-bit loads, or 64 bytes, which is lower than the maximum coalesced transaction size of 128 bytes 
1247      that can be issued by a warp. We can try recovering this performance by vectorizing the loads into <code>half2</code> loads like so:
1248    </p>
1249    <!-- <dt-code block language="clike"> -->
1250      <pre class="l-middle">        <code class="language-clike">
1251float2 sum01 = make_float2(0.0, 0.0);
1252for (int t = warp_id; t &lt; seq_len; t += t_stride) {
1253  float2 v01 = __half22float2(*((half2*)&amp;vh[kv_stride * t + i]));
1254  float att_t = atth[t];
1255  // Sadly CUDA does not have float2 SIMD ops
1256  sum01.x += v01.x * att_t;
1257  sum01.y += v01.y * att_t;
1258}
1259// atomicAdd both `sum01` lanes when we're done
1260        </code>
1261      </pre>
1262    <!-- </dt-code> -->
1263    <p>
1264      This attempt (<a href="https://github.com/andrewkchan/yalm/commit/ce8efcd20a5f1c04e69ff05a334cbc3aa951eac1">full patch here</a>) 
1265      improves the throughput somewhat, but it's still worse than the float32 version (188 μs vs. 140 μs, memory throughput 9.8% vs. 25.7%).
1266    </p>
1267    <p>
1268      What's going on? Our KV cache loads should be coalesced to the same sizes as the FP32 version. The final writes we do are 
1269      double the size, but reads far outnumber writes, and they should also still be coalesced. For hints, we collect a more 
1270      detailed <code>ncu</code> profile with <code>-set full</code> and <code>-o report</code> to get a file that we can view in the NCU UI.
1271    </p>
1272    <p>
1273      With this, we can see a very powerful line-by-line breakdown of the kernel being executed, with relevant stats and the corresponding 
1274      PTX and SASS instructions per line. In particular we see that the line loading from <code>vb</code> in the FP32 kernel is secretly 
1275      performing 8 32-bit loads:
1276    </p>
1277    <img src="./yalm-assets/source-cuda-kv32.png" class="l-middle" style="width: min(90vw, 700px)">
1278    <p>
1279      ...while the corresponding line in the FP16 kernel is performing only 1 32-bit load:
1280    </p>
1281    <img src="./yalm-assets/source-cuda-kv16.png" class="l-middle" style="width: min(90vw, 700px)">
1282    <p>
1283      Digging deeper, we can see that the compiled SASS instructions for the FP32 kernel have not only unrolled the loop into 4 iterations 
1284      (=8 loads) at a time, but are also arranging the load (<code>LDG.E</code>) and arithmetic (<code>FFMA</code>) instructions so that the loads can be 
1285      <a href="https://developer.nvidia.com/blog/boosting-application-performance-with-gpu-memory-prefetching/">prefetched</a>:
1286    </p>
1287    <img src="./yalm-assets/source-sass-kv32.png" class="l-middle" style="width: min(90vw, 700px)">
1288    <p class="caption">
1289      <i>
1290        The highlighted lines all correspond to the single selected line in the CUDA code above in the FP32 kernel.
1291      </i>
1292    </p>
1293    <p>
1294      Prefetching is a big boost for loops where each iteration loads from global memory and is independent from the results 
1295      of previous iterations. By unrolling 4 iterations of the loop and re-arranging all loads upfront, the compiler is able to hide 
1296      the latency of the loads for the latter 3 iterations.
1297    </p>
1298    <p>
1299      The FP16 kernel SASS does not even unroll any loop iterations, let alone rearrange loads:
1300    </p>
1301    <img src="./yalm-assets/source-sass-kv16.png" class="l-middle" style="width: min(90vw, 700px)">
1302    <p>
1303      So how do we get prefetching in the FP16 kernel? Unfortunately, just adding <code>#pragma unroll</code> to the loop does 
1304      not cut it (at least on my <code>nvcc</code> and machine setup). We are at the whim of compiler heuristics; they gaveth to our FP32 kernel 
1305      but now taketh away, and may continue to taketh away unless we instead 
1306      <a href="https://developer.nvidia.com/blog/boosting-application-performance-with-gpu-memory-prefetching/">manually unroll and prefetch values</a>. 
1307    </p>
1308    <p>
1309      Below, we have unrolled into 16 iterations (=32 loads) at a time, with some extra code to handle the iterations modulo 16. 
1310      We perform a batched prefetch<dt-fn>An even better solution might be to do a rolling prefetch using pipeline instructions. See the NVIDIA blog post linked above for more.</dt-fn> 
1311      where at the start of every 16th iteration, we issue a batch of loads, then retrieve them one-at-a-time until our batch is depleted. 
1312      The use of scalars instead of an array to hold prefetch results allows us to fetch to registers rather than slower 
1313      shared memory:
1314    </p>
1315    <!-- <dt-code block language="clike"> -->
1316      <pre class="l-middle">        <code class="language-clike">
1317float2 sum01 = make_float2(0.0, 0.0);
1318constexpr int UNROLL = 16;
1319half2 v01_0; float att_0; 
1320half2 v01_1; float att_1; 
1321half2 v01_2; float att_2; 
1322/* ... SNIP ... */
1323half2 v01_15; float att_15;
1324int t = warp_id;
1325for (int ctr = 0; ctr &lt; seq_len / t_stride - UNROLL + 1; t += t_stride, ctr++) {
1326  int ctr_mod = ctr % UNROLL;
1327  if (ctr_mod == 0) {
1328    // prefetch every UNROLL iterations
1329    #define PREFETCH(j) \
1330      v01_##j = *((half2*)&amp;vh[kv_stride * (t + j*t_stride) + i]); \
1331      att_##j = atth[t + j*t_stride];
1332    PREFETCH(0)
1333    PREFETCH(1)
1334    PREFETCH(2)
1335    /* ... SNIP ... */
1336    PREFETCH(15)
1337    #undef PREFETCH
1338  }
1339  // pull one value out of prefetch batch
1340  float2 v01;
1341  float att_t;
1342  switch (ctr_mod) {
1343    #define CASE(j) \
1344      case j: v01 = __half22float2(v01_##j); att_t = att_##j; break;
1345    CASE(0)
1346    CASE(1)
1347    CASE(2)
1348    /* ... SNIP ... */
1349    CASE(15)
1350    #undef CASE
1351  }
1352  // Sadly CUDA does not have float2 SIMD ops
1353  sum01.x += v01.x * att_t;
1354  sum01.y += v01.y * att_t;
1355}
1356// Handle any loop remainder that can't be unrolled
1357for (; t &lt; seq_len; t += t_stride) {
1358  float2 v01 = __half22float2(*((half2*)&amp;vh[kv_stride * t + i]));
1359  float att_t = atth[t];
1360  sum01.x += v01.x * att_t;
1361  sum01.y += v01.y * att_t;
1362}
1363// atomicAdd both `sum01` lanes when we're done
1364        </code>
1365      </pre>
1366    <!-- </dt-code> -->
1367    <p>
1368      The resulting kernel (<a href="https://github.com/andrewkchan/yalm/blob/5c7d634fdf88da76bd3d9ed98f517b49f8eb84d5/src/infer.cu#L393-L518">full code here</a>) finally gets us the expected wins, getting ~75 μs vs. ~140 μs for the FP32 kernel in our toy <code>ncu</code> benchmark. 
1369      And long context generations get faster as expected, from just over 57 tok/s to 58.8 tok/s. To top it all off, here's a video of 
1370      short context generation, where we get over 63.8 tok/s:
1371    </p>
1372    <div class="l-middle">
1373      <video controls="" src="./yalm-assets/gpu-kv16.mov" style="width: min(90vw, 720px)"></video>
1374    </div>
1375    <p>
1376      With this patch, we're now neck-and-neck with llama.cpp for long-context generations - at 58.7 tok/s vs. 58.8 tok/s - while handily 
1377      beating it for short contexts at 63.8 tok/s vs. 61.0 tok/s. Not bad!
1378    </p>
1379    
1380    <dt-byline></dt-byline>
1381    <h1 id="section-4">4. What's next</h1>
1382    <br>
1383    <p>
1384      Through this post, we've achieved close to state-of-the-art performance for a specific use case: local, completion-only, 
1385      single-batch, single-GPU inference on one model. We were able to do so with no libraries - just a few major optimizations 
1386      and a lot of CUDA/C++ elbow grease:
1387      </p><ul class="compact-li">
1388        <li>
1389          Multithreading and vectorization on the CPU
1390        </li>
1391        <li>
1392          Matmul warp reductions + coalescing, kernel fusing, and better attention in our GPU backend
1393        </li>
1394        <li>
1395          Weight + KV cache quantization for both backends
1396        </li>
1397      </ul>
1398    <p></p>
1399    <p>
1400      That said, we've only scratched the surface of LLM inference performance. Even in our specific use case there are lots of improvements to be made.
1401    </p>
1402    <p>
1403      One big opportunity is in the prompt pre-filling phase, where we have an existing sequence, similar to training time. A common optimization 
1404      here is to use matrix-matrix multiplication instead of matrix-vector to prefill multiple KV cache entries at once; this would improve the time-to-first-token,
1405      a metric just as important as raw token throughput for many applications. This also forms the basis of 
1406      the <a href="https://twitter.com/karpathy/status/1697318534555336961">speculative decoding</a> generation idea, which takes advantage of the fact that pre-filling K tokens is not that much slower than generating a single token.
1407    </p>
1408    <p>
1409      Another opportunity is to fuse our attention kernels, which the naive implementation split into separate scoring, softmax, and mixing operations. 
1410      This is essentially what FlashAttention does for training and FlashDecoding does for inference, with some clever insights about the underlying hardware.
1411    </p>
1412    <p>
1413      I'll also note that we are fast for FP16, but we're not FAST yet. Reaching 100 tok/s on our machine will require us to quantize more 
1414      aggressively, to smaller data types like FP8, INT8, or INT4, or to quantize activations<dt-fn>This is often done during training, 
1415      where we have large batches of activations that must be kept around for gradient computation in all layers. For single-batch inference it is less beneficial especially as we can reuse tensors between layers.</dt-fn>
1416      and KV cache entries as well as weights. Or to do computations natively in lower precision!
1417    </p>
1418    <p>
1419      At a low level, there are lots of kernel-specific implementations to be made. Many of the kernels we ended up with still display 
1420      suboptimal performance, with several having memory throughput of less than 60% of the speed-of-light. Kernel optimization is a 
1421      huge topic and many of the optimizations we discussed - memory coalescing, vectorization, pre-fetching, etc. - need to be combined 
1422      in complex ways to really reach peak performance.
1423    </p>
1424    <p>
1425      Finally, in a production implementation, it's much more practical to use libraries of kernels maintained by experts, which are usually 
1426      state-of-the-art, like cuDNN kernels and cuBLAS. These libraries have had decades of work go into them and are made by the same folks designing 
1427      the chips themselves, to the point where even approaching <a href="https://siboehm.com/articles/22/CUDA-MMM">within 5% performance of cuBLAS</a> 
1428      on specific hardware and use case is considered an achievement. Of course, this route is also usually less fun :)
1429    </p>
1430  </dt-article>
1431
1432  <dt-appendix>
1433    <p>
1434      Thanks to Thariq Shihipar, Jonathan Chan, and Luciano Vinas for reading and giving feedback.
1435    </p>
1436    <h3 id="appendix">Appendix</h3>
1437    <p>
1438      Following is the code used to benchmark huggingface transformers. Note the benchmarked program does not use <code>torch.compile</code>; my 
1439      attempts to do so following the guide in the <a href="https://huggingface.co/docs/transformers/en/llm_optims#static-kv-cache-and-torchcompile">
1440      HuggingFace LLM optimization docs</a> either did nothing or resulted in worse throughput, despite HuggingFace transformers supposedly supporting
1441      the <a href="https://github.com/huggingface/transformers/issues/28981">use of <code>torch.compile</code> with Mistral</a>.
1442    </p>
1443    <!-- <dt-code block language="python"> -->
1444      <pre>        <code class="language-python">
1445import time
1446
1447from transformers import AutoModelForCausalLM, AutoTokenizer
1448import torch
1449
1450mistral_models_path = "../Mistral-7B-Instruct-v0.2"
1451
1452tokenizer = AutoTokenizer.from_pretrained(mistral_models_path)
1453
1454model = AutoModelForCausalLM.from_pretrained(mistral_models_path, torch_dtype=torch.float16)
1455model.to("cuda")
1456
1457# short prompt
1458model_inputs = tokenizer(["Q: What is the meaning of life?"], return_tensors="pt").to("cuda")
1459
1460# do a warmup inference
1461generated_ids = model.generate(**model_inputs, max_new_tokens=1, do_sample=True)[0].tolist()
1462
1463# benchmark
1464start_time = time.time() 
1465generated_ids = model.generate(**model_inputs, max_new_tokens=1000, do_sample=True)[0].tolist()
1466# decode with mistral tokenizer
1467result = tokenizer.decode(generated_ids)
1468end_time = time.time()
1469
1470elapsed_s = end_time - start_time
1471print(result)
1472
1473num_prompt_tokens = model_inputs["input_ids"].shape[-1]
1474num_generated_tokens = len(generated_ids) - num_prompt_tokens
1475
1476print(f"Generation stats:\n" +
1477      f"  prompt: {num_prompt_tokens} tokens\n" +
1478      f"  generated: {num_generated_tokens} tokens\n" +
1479      f"  throughput: {num_generated_tokens/elapsed_s}tok/s\n" +
1480      f"  latency: {elapsed_s/num_generated_tokens}s/tok\n" +
1481      f"  total: {elapsed_s}s\n")
1482        </code>
1483    </pre>
1484  <!-- </dt-code> -->
1485  </dt-appendix>
1486
1487  
1487<script type="text/bibliography">
1488  </script>
1488
1489  
1489<script defer="" src="https://cdn.jsdelivr.net/npm/[email protected]/dist/contrib/auto-render.min.js" integrity="sha384-kmZOZB5ObwgQnS/DuDg6TScgOiWWBiVt0plIRkZCmE6rDZGrEOQeHM5PcHi+nyqe" crossorigin="anonymous"></script>
1489
1490  
1490<script>
1491    document.addEventListener("DOMContentLoaded", function() {
1492      renderMathInElement(document.body);
1493    });
1494    const toc = document.getElementById("toc");
1495    const addItemToTOC = (h, title, id, indentLevel) => {
1496      const a = document.createElement("a");
1497      a.href = `#${id}`;
1498      a.textContent = title;
1499      const li = document.createElement("li");
1500      li.classList.add(`toc-${h}`);
1501      for (let i = 0; i < indentLevel; i++) {
1502        const tab = document.createElement("span");
1503        tab.innerHTML = "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;";
1504        li.appendChild(tab);
1505      }
1506      li.appendChild(a);
1507      toc.appendChild(li);
1508    }
1509    addItemToTOC("anchor", "(Top)", "", 0);
1510    // Add `<li><a href="#section-...">{sectionTitle}</a></li>` to TOC for all sections
1511    const sections = document.querySelectorAll(`[id^=section-]`);
1512    for (let section of sections) {
1513      const h = section.tagName;
1514      const title = section.textContent;
1515      const id = section.id;
1516      let indentLevel = 0;
1517      for (let i = 0; i < id.length; i++) {
1518          if (id[i] === ".") {
1519              indentLevel++;
1520          }
1521      }
1522      addItemToTOC(h, title, id, indentLevel);
1523    }
1524    addItemToTOC("anchor", "(Comments)", "bottom", 0);
1525  </script>
1525
1526</article>
1527  <div id="bottom"></div>
1528  
1528<script src="https://utteranc.es/client.js" repo="andrewkchan/andrewkchan.github.io" issue-term="pathname" theme="github-light" crossorigin="anonymous" async="">
1529  </script>
1529
1530
1531</body></html>

Line numbers count LF bytes from the start of the resource, as the search results do. Vendor segments are library code the classifier recognised; they are stored but not indexed. Bytes are shown as Latin1 characters, one per byte.