LLM高速化勉強会資料

Transcript

1. LLM高速化(勉強会) LLM推論システム勉強会の方がよかった気がしなくもない
2. 目次 ・自己回帰型モデルでの言語生成方法のおさらい ・実際に推論してみよう(通常のtorch) ・アルゴリズム的な高速要素: KVCache, FlashAttention, Super Sequence + Continual Batching, PagedAttention, Speculative Decoding ・アルゴリズム的な高速要素の実装: CUDA, Triton, CuTe, KernelFusion ・アーキテクチャ面: GQA, MLA, Sliding Attention, Linear Attention ・量子化 ・Profiler: Nsight Compute, Nsight Systems + NVTX ・CUDAGraph ・vLLMの紹介とコントリビュート方法: Chunked Prefill, Preemption, Prefill-Decode Disaggregation ・実際に推論してみよう(vLLM) ・複数ノード分散(概念のみ): MPI, Data Parallel, Pipeline Parallel, Tensor Parallel, Context Parallel ・NVIDIA製GPU: Tensor Core, TMA, Compute Capability
3. 注意 ・LLM Servingを支える技術 という記事にInspiredされて始めたイベントです ・上の記事に個人的な経験も含めて発表を行いますが, 間違っていたら上の記事に書い てあることの方がおそらく正しいので指摘してください
4. 自己回帰型モデルでの言語生成方法のおさらい: 言語モデル
5. 自己回帰型モデルでの言語生成方法のおさらい: Transformer Decoder
6. 自己回帰型モデルでの言語生成方法のおさらい: Attention
7. 本日のまとめ Attention部分: ・KVCacheでDecodeを高速化 ・PagedAttentionでメモリ使用量・ アクセス量を効率化 ・GQAなどのアーキテクチャで計 算量とKVCacheメモリを削減 Tokens生成部分: ・Speculative Decodingで1度に予測す るTokens数を増加 処理全体: ・量子化でメモリ使用量 , 転送量 を削減 ・CUDAGraphでCPU待ち時間を 削減 ・Super Sequence + Continual Batchingで余分なメモリ使用量を 削減 ・PrefillとDecodeは分離するか同 時に処理するかしてリクエストの 割り込みを許可 ・大きいモデルは並列分散処理
8. まずはともあれ推論してみよう URL: https://colab.research.google.com/drive/1ZVHxBKH9L4r3yf5AF7hEzMP7CONHLb 9l?usp=sharing 今回のイベントを通して使用するモデルはQwen/Qwen2.5-0.5B-Instruct, Colabの無 料枠で使えるTesla GPUのハードウェア上で動くモデルかつ軽量なもので適合したのが これであった Transformersの部分までの推論を動かすと, 現時点では5.97 tokens / secとなるはず 後半のvLLMは94.84 tokens / secとなるので最後まで聞くと何故Transformersを使っ た推論が遅いのかという問題に答えられるようになるはず
9. アルゴリズム的な要素での高速化: KVCache ・文章を生成する時に必要な情報 はAttention計算後の最後の Tokensが持つ隠れ層 ・この最後の隠れ層はKとVの計測 結果をCacheし続けておくことで計 算量を減らすことが可能(右図) ・余談として最初のPrompt全てを 計算してKVCacheを構築する段階 をPrefill, KVCacheを使用して1トー クンずつ予測していく段階を Decodeといいます。
10. ・正直説明されてもよくわからないと思うのでColabを作成しました https://colab.research.google.com/drive/1uZLXqNGUsoRwFOGQQ8pQ422iNQvU mEP9?usp=sharing ・KVCacheが長系列の推論時にはメモリボトルネックになりえます。KVCacheのsizeは 以下の式で求められます 2 * batch_size * num_layers * sequence_length * head_dim * num_heads ・batch_size = 16, num_layers = 96, sequence_length = 32768, head_dim = 128, num_heads = 8, float16とかで計算をしてみるとなんと192GiBとなります ・KVCacheのsizeを減らすための研究は後々話します アルゴリズム的な要素での高速化: KVCache
11. アルゴリズム的な要素での高速化: FlashAttention ・前のスライドで紹介したColabではQ, K, V, scores, weights, headsなどが計算の 度にプロセッサのメモリから取得され, 書き 戻されている
12. アルゴリズム的な要素での高速化: FlashAttention ・GPUのメモリはHBMを使用しておりTB/s 級の帯域幅を持つものの, 依然として複数 回大きなデータを書き込むとボトルネックに なりうる ・GPUのHBMとSRAMは右図によると 12.66倍ほどの帯域幅の差がある, どうに かしてこのSRAMをなるべく頻繁に使えな いか
13. アルゴリズム的な要素での高速化: FlashAttention ・FlashAttentionではTiling QKの計算とOnline Softmaxの計算でHBMアクセスを少なく しつつなるべくデータをSRAMに乗っけて計算している
14. Super Sequence + Continual Batching ・PagedAttentionの前にvLLMなどで使われているBatching方法について解説する ・推論システムとしては複数のリクエストをバッチ処理する必要がある ・この時PyTorchでやるようにBatchを組み立てるとRequestのtoken_idの列は (batch_size, max_n_tokens)となるため, リクエストの中で最もtokens長が長い部分に 引っ張られ, 余分なpaddingが生じてしまう ・Super Sequence (命名はvLLMのブログより)ではリクエストのtokens列を2次元のテ ンソルではなく1次元のtokens列にして処理を行い,paddingを無くしている
15. Super Sequence + Continual Batching Super Sequenceにより新しくリクエストがきた時には末尾に新しいtokens列をくっつけ るだけでよくなる。現在きているリクエストを処理しながら新しいリクエストを受け付けるこ の方法をContinual Batchingという その場合, Attentionのmask計算を工夫することで違いの文章の相互作用を防いでいる (次スライド)
16. Super Sequence + Continual Batching
17. アルゴリズム的な要素での高速化: PagedAttention PagedAttention はOSのPaging機構からインスピレーションを受け, KVCacheをBlock 単位で管理し, Batch推論におけるKVCacheの余分な断片化, Paddingを無くす手法で ある。 詳しくはvLLMのブログのgifが詳しいため確認するが, 実際には概念だけではなく vLLM内部でこれを可能にする実装が偉い(HPCは実装の方が大切)
18. アルゴリズム的な要素での高速化: Speculative Decoding 通常のLLM推論ではK tokensの予測にK回のforwardが必要になるが100B以上級の モデルをK回forwardするとlatencyは無視できなくなる。 出力する確率分布の近い軽いモデルや複数tokensの予測を行うmoduleを用意し, 仮 tokensを予測させて大きなモデルに推論させ, 棄却サンプリングで複数tokensを採択さ せる手法がSpeculative Decodingである
19. アルゴリズム的な要素での高速化: Speculative Decoding ・最も簡単なSpeculative Decodingは小さいモデルに複数Tokensを予測させ, 大きなモ デルで棄却サンプリングをする方法である ・しかしこの方法だと小さいモデルと小さいモデル自体のKVCacheも膨れ上がるため実 用性が低い
20. アルゴリズム的な要素での高速化: Speculative Decoding ・そのため, 数tokens先をなるべく余分なメモリを使わずにDraftする手法がいくつか提 案されている。MedusaとEagleはそのうちの1つ。図はここから拝借
21. アルゴリズム的な要素での高速化: Speculative Decoding Gemma4のMTPがやや私が知らないものだったので少し驚きました。みんなで見てい きましょう ・いずれにせよやはりKVCacheを共有していることからもKVCacheのせいでメモリをあ まり多く使いたくないという技術選定はあるかと思います(意見)
22. アルゴリズム的な要素での高速化: Speculative Decoding ・最後の2段階のCluster行列の計算を見積もってみます。 通常の1 tokensあたりの行列積の計算量はhidden_dim * vocab_sizeです。 ・2段階のCluster行列の場合計算量はhidden_dim * (cluster_size + vocab_size_in_cluster)です。
23. アルゴリズム的な要素での高速化: Speculative Decoding 均等に割り振られているとしてvocab_size_in_cluster = vocab_size / cluster_sizeとし てみます。 1 <= cluster_size <= vocab_sizeであることに注意すると(相加相乗平均などで整理), 大雑把に評価して 2√vocab_size <= cluster_size + vocab_size / cluster_size <= vocab_size + 1 となるので, 大体の極端じゃないcluster_sizeにおいては多分cluster_size + vocab_size / cluster_sizeはvocab_sizeより小さくなり, Cluster行列にした方が早いと 思います(あと2段階目のメモリloadが減るのも嬉しい) 余談: 等号成立条件は工学部なので確認しない
24. アルゴリズム的な要素での高速化の実装 ・先のスライドでFlashAttention, PagedAttentionなどを紹介した ・単純にPython + Torchで実装するとPythonのfor文が実行するだけで10~100μs消費 してしまうなどの問題があるため, 基本的にGPUで動かさなければならない ・実装のツールとしてCUDAとTriton, CuTeを紹介する
25. アルゴリズム的な要素での高速化の実装: CUDA ・GPUで動くコードを書く上で直接アセンブリを書く次に原始的な方法 ・話すと長すぎるので省略する ・LeimaoさんのCUDA Matrix Multiplication Optimizationで一回はCUDAを書くといい と思います
26. アルゴリズム的な要素での高速化の実装: CUDA ・CUDAの最適化の80%はメモリをどう効率的に読んでいくかを最適化することになる ・CUDAはshared memoryと呼ばれるConfigurableなCache領域を持つ, ここに複数回 アクセスされるデータをloadして最適化することが多い ・余談: このような高速な有限メモリと低速な大容量メモリの間を扱うような理論として Red-Blue Pebble gameとかCommunication Avoiding Algorithmがあるので調べてみ てください
27. ・CUDAを書くのは工数がかかり, 専門のエンジニアを雇うコストがかかる ・そこでPythonからでも高性能なCUDA Kernelがかけることを目指したDSLがTritonと CuTe-DSL である ・ここ最近のFlashAttentionやvLLMも基本TritonかCuTeでかかれている印象がある ・余談: SC25というスパコンのカンファに参加した時もNVIDIAは基本的にPythonからか けるものをどんどん提供していくというスタンスだった アルゴリズム的な要素での高速化の実装: Triton
28. アーキテクチャ面: GQA, MLA ・普通のMulti-Head Attentionは1つのQueryに対して1 pairのKVが対応したが Grouped Query Attentionでは複数のGrouped Queryに対して1 pairのKVが対応する ・これにより1トークンあたりのKVCacheの量が少なくなる
29. アーキテクチャ面: GQA, MLA ・Multi-Latent AttentionはKVのhidden_dimを縮小して拡大したものを計算に使う Attentionである ・hidden_dimを縮小したKVをKVCacheすることでKVCacheのメモリを減らしている
30. Multi-Head Latent Attention: 具体的な計算 (昔計算した) ・low-rankにすることの意義を計算する。計算に使う parameterはGitHubから取得, 計算は1 batch分で行う ・従来はSequence × hidden_dim -> n_heads * (k_head_dim + q_head_dim)の計算が走る ・なので計算量は Sequence(S) × hidden_dim × n_heads × (k_head_dim + q_head_dim) = 8,388,608 × S, KV Cacheは4096×F×S bytesだけ保存する(Fはfloat32で4, float16で2) ・ここで512にlow-rankにすると ・計算量はS × hidden_dim × low-rank + S × low-rank × n_heads × (k_head_dim + q_head_dim) = 3,145,728 × S, KV Cacheは512×F×S bytesだけ保存する ・以上より, 計算量は3 / 8に, KV Cacheの容量は1 / 8になる →AttentionのKV部分の計算時間は 3 / 8に, KV Cacheのload時間が1 / 8に, しかも同一容量で 8倍の sequenceを扱うこともできる ・このあたりは性能とメモリ削減量のトレードオフを見ながら調整されているだろうことを考えると相当の努力が 見られる
31. アーキテクチャ面: Sliding Attention ・一部のLayerに関してKとVは過去window size分だけ見ればよくない?というAttention, KVCacheのmaxが固定値になる点が嬉しい ・Gemmaが基本このアーキテクチャである
32. アーキテクチャ面: Linear Attention ・Qwen3.5以降のQwenで採用されているAttention ・Attention部分の計算が小さく, KVCacheのメモリ消費量も小さい ・おそらく実装はここ https://github.com/fla-org/flash-linear-attention/blob/6f62104c1ec677e5951ac5464 2fd166e765db54e/fla/ops/common/fused_recurrent.py#L30 ・KVの代わりに前のhidden_stateを使う都合上, 採択tokens数に応じてSpeculative Decodingのような途中までのKVCacheを廃棄するような手法に弱い
33. 量子化 パラメーターを表現するための必要なbit数を減らすことでメモリ容量を削減し, かつ計算 時のメモリストールの時間を減らすことができる技術 GPT-OSSなどのMoE部分はFP4で量子化されていたりする。これはMoEが全体のうち 90%以上を占めるメモリボトルネックであるためである 精度を保つための方法として以下の工夫がある ・Block-wise Scaling ・Mixed Precision Accumulation Mixed Precisionに関して少し面白い話があるので紹介する
34. 量子化: Mixed Precision Accumulation ・一部を量子化し, 一部をfloat16やfloat32で演算することでメモリ使用量と精度を保つ 方法 ・積和演算は誤差が蓄積するため, GPU内ではFP16もしくはFP32で演算することで精 度と速度とメモリ使用量を保っている(下図はDeepSeek-R1のFP8 GEMM Kernel)
35. 量子化: Mixed Precision Accumulation ・多分HPCの人間にしか知られていない手法として, Ozaki Schemeが存在する ・Ozaki Schemeは低精度計算を利用して高精度の行列演算をEmulateする手法であ る ・つまり, FP32の精度を計算するのにFP16のTensor coreやハードウェアを使うことが でき, FP32演算のカタログ値以上の速度を出すことが可能となる ・これはcuBLASにも統合されている
36. Profiler: Nsight Compute ・CUDA Kernel 1つの性能を検証するためのツール ・遅延がメモリストールによる遅れなのかSM(GPUのコアのようなもの)の飽和によるも のなのかなどを確認するのに使う
37. Profiler: Nsight Systems ・NVIDIA GPUを使ったシステム全体での計測を行うツールでボトルネックの特定に使う ・NVTXというツールでコード中のどの場所に対する測定なのかを可視化してくれ , 便利
38. CUDAGraph ・Nsight Systemsで最初の方に実行した推論のProfileを取るとGPUのKernelとKernel の間に100μs程度の間が空いてしまっているタイムラインがよく見える ・実は何も最適化されていない場合のtorchでの推論ではこれが最も大きいボトルネック である。 ・何が起きているか: GPUの実行速度にCPUが追いつけていない
39. CUDAGraph ・これを解決するのがCUDAGraphという機能 ・CUDA Graph は、多数のカーネル起動やメモリ転送といった操作をあらかじめ「グラ フ」として記述し、その後まとめて実行・再利用できる仕組みです
40. CUDAGraph ・GPU に対するカーネル呼び出しやデータコピーはそれぞれホスト側から API を呼び 出すたびにオーバーヘッドが発生しますが一度このグラフをキャプチャしておけば全工 程をまとめて走らせられます ・カーネル起動やメモリコピーのオーバーヘッド、CPUでのPythonに起因するオーバー ヘッドなどを大幅に削減できます
41. vLLM ・ここまで紹介してきた手法をまとめて1つの推論エンジンソフトウェアとしたもの
42. vLLM: アーキテクチャ 全体のアーキテクチャ: ユーザーからのリクエストを受けるFrontend部分とEngineCore 部分に分かれ, Engine Coreがリクエストをスケジューリングするスケジューラーと, KVCacheのmanageを担当するKVCacheManager, 実際の計算を実行するModel Executorを持っている
43. vLLM: Chunked Prefill vLLMはSuper SequenceによりPrefillとDecodeを区別なく扱うことができる。 途中で割り込んだPrefillが全体の処理時間を圧迫しないように, Prefillを分割して行って いる。デフォルトではFalse
44. vLLM: Preempt vLLMは実運用に耐えるためにメモリがOverしたときの処理も書いてある KVCacheの物理ブロックが尽きた時に一部のリクエストをPreemptし, KV ブロックを解 放して他のリクエストに回す。 中断の仕方として, 復活時に再計算するかCPUにKVCacheをOffloadするかがある。 基本的に後からきた順でPreemptを行う Scheduler Classが担当を行う。
45. vLLM: Prefill-Decode Disaggregation PrefillでとDecodeでは計算の特性の違いからPrefillとDecodeを混在させるとシステム 全体のパフォーマンスが下がりうる そこでPrefillを行うサーバーとDecodeを行うサーバーを分離することでTime To First TokenとTime Per Output Tokenを個別に最適化が可能となる
46. vLLM: Prefill-Decode Disaggregation また, 計算の特性からそれぞれの段階で使うハードウェアの調整も可能となる 例: Prefillは計算優位なのでTensor coreの強いH200, DecodeはA100程度でも十分な速 度になるなど
47. vLLM: Contribute方法 vLLMは改造するのに優しいソフトウェアです❗ ・sourceからのビルド方法が書いてある ・全体アーキテクチャの説明も丁寧 ・簡単な実験だけならPluginsでもできます！ 俺もやってるんだからみんなもやってください ・強いて言うなら何もしないとMultiprocessになってしまいDebuggerで追うのが困難に なるのでos.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"で Inprocessにしておくべきです
48. まとめ Attention部分: ・KVCacheでDecodeを高速化 ・PagedAttentionでメモリ使用量・ アクセス量を効率化 ・GQAなどのアーキテクチャで計 算量とKVCacheメモリを削減 Tokens生成部分: ・Speculative Decodingで1度に予測す るTokens数を増加 処理全体: ・量子化でメモリ使用量 , 転送量 を削減 ・CUDAGraphでCPU待ち時間を 削減 ・Super Sequence + Continual Batchingで余分なメモリ使用量を 削減 ・PrefillとDecodeは分離するか同 時に処理するかしてリクエストの 割り込みを許可 ・大きいモデルは並列分散処理
49. もう一回ハンズオンしてみよう https://colab.research.google.com/drive/1ZVHxBKH9L4r3yf5AF7hEzMP7CONHLb 9l?usp=sharing ・5.97->94.84で15.88倍❗
50. 他の話題 ・Roofline modelは見てください, 大体は行列積ぐらいじゃないとCompute boundになら ないです ・複数ノード分散(概念のみ): MPI, Data Parallel, Pipeline Parallel, Tensor Parallel, Context Parallel ・NVIDIA製GPU: Tensor Core, TMA, Compute Capability