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& model, std::string prompt, int steps) { 206 std::vector<int> 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 < 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 << tokenizer.decode_one(next_token) << 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& 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->token_embedding_table); 252 253 // Models consist of a sequence of transformer blocks which 254 // mutate the inference state in order 255 for (Block& b : this->blocks) { 256 b->block(s); 257 } 258 259 // Usually there is a layer norm right before the final classifier 260 s.x = layernorm(s.x, this->lm_head_prenorm_weights); 261 // Typically we end with a linear transform from (dim) -> (vocab_size) 262 s.logits = linear(s.x, this->lm_head_classifier_weights); 263} 264 265void Block::block(InferenceState& s) { 266 s.x_resid = layernorm(s.x, this->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->wq,
275 this->wk, 276 this->wv, 277 this->key_cache, 278 this->value_cache 279 ); 280 s.x += s.x_resid; 281 s.x_resid = layernorm(s.x, this->ffn_prenorm_weights); 282 // On modern architectures like Llama, this is a GLU feedforward 283 // with 3 linear transforms, not a simple MLP: 284 // -> 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->w1, this->w2, this->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,) -> xout (d,) 414 int i; 415#pragma omp parallel for private(i) 416 for (i = 0; i < d; i++) { 417 float val = 0.0f; 418 for (int j = 0; j < 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__) && defined(__F16C__) 484 // W (d,n) @ x (n,) -> xout (d,) 485 assert(n % 16 == 0); 486 int i; 487#pragma omp parallel for private(i) 488 for (i = 0; i < 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 < 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*)&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(&x[j]); 502 __m256 xvechi = _mm256_loadu_ps(&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 && "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><<< numBlocks, threadsPerBlock >>></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><<< . >>></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<<<c.hidden_dim, WARP_SIZE>>>( 567 w1<float>(), s.xb(), c.dim, c.hidden_dim, s.hb() 568); 569matmul<<<c.hidden_dim, WARP_SIZE>>>( 570 w3<float>(), s.xb(), c.dim, c.hidden_dim, s.hb2() 571); 572glu_gelu<<< 573 (c.hidden_dim + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK, 574 MAX_THREADS_PER_BLOCK, 575>>>( 576 s.hb(), s.hb2(), s.hb() 577); 578matmul<<<c.dim, WARP_SIZE>>>( 579 w2<float>(), s.hb(), c.hidden_dim, c.dim, s.xb2() 580); 581// ffn residual back into x 582add_residuals<<< 583 (c.dim + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK, 584 MAX_THREADS_PER_BLOCK 585>>>( 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<T>(), c.dim, c.hidden_dim); 599matmul(s.hb2(), s.xb(), w3<T>(), c.dim, c.hidden_dim); 600for (int i = 0; i < c.hidden_dim; ++i) { 601 s.hb()[i] = gelu(s.hb()[i]) * s.hb2()[i]; 602} 603matmul(s.xb2(), s.hb(), w2<T>(), c.hidden_dim, c.dim); 604// residual connection back into x 605for (int i = 0; i < 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,) -> out (d,) 629 int i = blockIdx.x * blockDim.x + threadIdx.x; 630 if (i >= d) return; 631 float sum = 0.0; 632 for (int j = 0; j < 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<<< 641 (d + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK, 642 MAX_THREADS_PER_BLOCK 643>>>(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 > 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 < 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,) -> out (d,) 687 // PRECOND: Blocks are 1-D and same size as warp. 688 int i = blockIdx.x; 689 if (i >= d) return; 690 int offset = threadIdx.x; 691 float rowSum = matmul_row(&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<<<d, BLOCK_SIZE>>>(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 < blockDim.x / warpSize ? sm[lane] : def; 814} 815 816template <typename T> 817__global__ 818void matmul_wide(const T* A, const float* x, int n, int d, float* out) { 819 // A (d,n) @ x (n,) -> 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 >= d) return; 823 // Warp j computes sum for row at <blockIdx.x*WPB + j> 824 // Lane 0 of each warp will hold result 825 int k = threadIdx.x % warpSize; 826 float rowSum = matmul_row(&A[n * i], x, k, n); 827 // Transpose values so lane k in warp 0 contains row at <blockIdx.x*WPB + k> 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 < 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 >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 > 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 < 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<<<blocks, tpb>>>( 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 < head_dim; i += blockDim.x) { 1063 float sum = 0.0; 1064 for (int t = t_start; t < t_start + t_per_thread; t++) { 1065 sum += vh[kv_stride * t + i] * atth[t]; 1066 } 1067 atomicAdd(&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<<<blocks, tpb>>>( 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 < 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 < seq_len; t += t_stride) { 1139 sum += vh[kv_stride * t + i] * atth[t]; 1140 } 1141 atomicAdd(&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<<<blocks, tpb>>>( 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 < 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 < 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 < seq_len; t += t_stride) {
1253 float2 v01 = __half22float2(*((half2*)&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 < 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*)&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 < seq_len; t += t_stride) {
1358 float2 v01 = __half22float2(*((half2*)&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 = " "; 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.