vLLM の中身を覗いて、自作 kernel で LLM の生成速度を約15%上げた話

はじめに — 「もっと速く、もっと深く」最適化したい

これまで LLM のサービング（学習済みモデルに文章を生成させ、サービスとして提供する仕組み）について、何本か記事を書いてきました。次に挑戦したくなったのは、フレームワークの設定やパラメータを調整する段階から一歩踏み込んで、GPU kernel のレベルまで降り、より最先端の手法で推論を最適化することです。「より速く、より深く」——それがこの記事のモチベーションです。

ローカルで LLM を動かす選択肢と vLLM

ローカルで LLM を動かすだけなら、LM Studio や Ollama のように、GUI やコマンド一発で手軽に使えるツールがあります。まず試すならこれで十分です。

一方で、もっと多くのリクエストを高いスループットで捌いたり、本番のサービングに近い形で性能を突き詰めたいときの定番が vLLM です。vLLM は LLM 推論に特化した「推論エンジン」で、PagedAttention（GPU メモリ上の KV cache を OS のページングのように効率管理する仕組み）や continuous batching（リクエストを動的に束ねて GPU を遊ばせない仕組み）といった工夫により、GPU を高い使用率で回しながら高スループット・低レイテンシを狙えます。研究でも本番でも広く使われ、最新モデルへの対応も速い定番の一つです。導入も pip install vllm で入り、今回のローカル GPU（RTX 4070）でもすぐ動かせます。

もし vLLM を触ったことがなくても大丈夫です。この記事は「推論エンジンの中身を計測して、自作 kernel で速くしていく話」として読めます。

この記事でやること

そして vLLM は OSS です。実装がすべて公開されているので、用途に合わせて自分で手を入れることもできます。ここに、最先端の kernel 最適化を実地で試す余地がありました。

実際、あるモデル（今回は Qwen3.5-2B）で vLLM の decode（1 トークンずつ文章を生成していく処理）を評価していくうちに、実装の中に まだ最適化しきれていない箇所 を見つけました。そこを自作 kernel に置き換えてみると、なんと vLLM 上での 1 トークンあたりの生成にかかる時間が 約 15% も短くなった のです。

この記事は、その vLLM の内側を実際に触って速度を上げた過程 の記録です。プロファイラ（Nsight）で「どこが遅いか / どこに無駄があるか」を探し、狙いを定めて処理を自作 kernel に置き換え、実際に効果を確かめる——という一連の流れをまとめました。

要点（結論）

RTX 4070 で vLLM + Qwen3.5-2B の文章生成の中身を Nsight で覗いてみると、行列計算（GEMV。1 トークンずつ生成する decode では、行列 × ベクトルの計算になります）そのものではなく、その周りに残っていた **GemmaRMSNorm（各層で数値の大きさを整える正規化処理）**が、細かい kernel にバラバラに分かれて動いていました。ここを自作の Triton kernel 1 つにまとめただけで——

1 トークンあたりの生成時間（TPOT）が、約 15% 短くなりました。

max tokens	TPOT（改善前 → 改善後）	decode tokens/s	総レイテンシ削減
128	10.615 → 8.919 ms	94.214 → 112.129	15.88%
512	10.557 → 8.927 ms	94.720 → 112.017	15.43%
2048	10.625 → 9.037 ms	94.121 → 110.659	14.94%

計測条件は Qwen3.5-2B / vLLM nightly / --enforce-eager、短い prompt から長めに生成させる decode 中心の workload です。この数字はこの条件でのものなので、どんな設定でも同じだけ速くなるわけではありません（詳しくは後半の「限界」で）。

伝えたいことは 1 つです。自作 kernel の価値は、単体ベンチマークの速さだけでは決まらない。 実際に動いている推論エンジンを計測して、既存実装に残った「無駄」を見つけ、そこを狙って書き換え、最後に生成速度がちゃんと上がったかまで確かめる——その一連の流れにこそ価値がある、ということです。

リポジトリ: https://github.com/CHIPMUNK-T0T/cuda-kernel-engineering

対象読者・検証環境（クリックで展開）

対象読者

• ローカル LLM に興味があり、vLLM のような推論エンジンの「中身」を覗いてみたい人
• 「GPU kernel を自作して速くする」という作業が具体的にどういう流れなのか知りたい人
• LLM 推論 backend をプロファイルして、遅い / 無駄な箇所を見つけたい人
• 自作 CUDA / Triton kernel が「実 backend で本当に効くのか」を知りたい人

前提知識は多くなくても読めるように、用語は本文で都度補足します。CUDA / PyTorch を少し触ったことがあると、より楽しめると思います。

検証環境

- GPU: RTX 4070 (12GB)
- CUDA: 12.8 / PyTorch 2.11.0+cu128 / Triton 3.6.0 / Python 3.12
- backend: vLLM nightly, --enforce-eager
- model: Qwen/Qwen3.5-2B (bfloat16)
- workload: short prompt + decode-heavy, max_tokens = 128 / 512 / 2048
- 計測: stream=True は warmup=1, runs=3。非 streaming の同条件比較は warmup=3, runs=5
- profiler: Nsight Systems (request-only trace), Nsight Compute (kernel 単体)

--enforce-eager という実行モード

本題に入る前に、計測で使った --enforce-eager について少しだけ説明させてください。

LLM を動かすとき、vLLM は普段、生成をできるだけ速くするために内部でいくつもの高速化の仕組みを使っています。代表的なのが CUDA Graph や torch.compile です。ざっくり言うと、たくさんの細かい処理を「ひとかたまり」にまとめて GPU に流し、余計な待ち時間を減らす仕組みです。実行は速くなる一方で、中で実際に何がどう動いているかは計測ツールから見えにくくなります。小さな処理がまとめられて 1 つの塊になってしまうので、「ここに無駄がある」と指差すのが難しいのです。

--enforce-eager は、その「まとめる高速化」をあえてオフにして、処理を 1 つずつ素直に実行させるモードです。速さの面では少し不利になりますが、そのぶん内部の処理が、プロファイラ（Nsight）のタイムラインにそのまま並んで見えるようになります。どの処理が、何回、どれくらいの時間動いているのか——これがはっきり分かる状態になるわけです。

なぜ中を覗くことになったのか

「いつか kernel のレベルまで降りて最適化してみたい」という気持ちは、前々からありました。ただ、正直に言うと、この調査も最初から「よし、--enforce-eager で中を観察して最適化ポイントを探すぞ」と意気込んで始めたわけではありません。着手のきっかけは、もっと素朴な事情でした。

長いコンテキストで動かそうとすると、RTX 4070 の 12GB のメモリでは足りなくなることがあります。先ほどの高速化の仕組みはその分メモリも余分に使うため、少しでも節約しようと、半ば仕方なく --enforce-eager を付けて動かしていました。

その状態でモデルを回しているうちに、ふと「これ、中では実際にどんな計算をしているんだろう？」と気になって、プロファイラで覗いてみたのが始まりです。LLM の推論には、大きく分けて、入力した prompt をまとめて処理する prefill と、そのあと 1 トークンずつ文章を生成していく decode の 2 段階があります。この prefill と decode で、どの処理がどれくらいの時間を食っているのかを眺めていました。

すると decode の側で、意外なことに気づきます。本来いちばん重いはずの行列計算そのものだけでなく、その周りにある正規化（GemmaRMSNorm）のような小さな処理が、細かく分かれたまま何度も動いていて、無視できない時間を占めていたのです。この「本命ではないところに残った無駄」が、今回、自作 kernel で潰しにいった相手でした。

いろいろ計測して分かったこと

ここから対象を絞り込むまでに、いくつも計測を重ねました。細かい表や数字は割愛して、分かったことだけをざっくりまとめます。

• 重い行列計算そのものは、自分では速くできませんでした。 ここは cuBLAS という NVIDIA 純正の実装がすでに非常に強く、個人が正面から書き換えても勝てません。
• 効きそうだったのは、その周りの小さな処理でした。 コピーや型変換、正規化といった処理は 1 回あたり 1 マイクロ秒ほどですが、decode では 1 トークンごとに何度も呼ばれるので、積み重なると無視できない時間になります。
• ただし「小さい処理を単純に 1 つにまとめる」だけでは、かえって遅くなることもありました。まとめて効くのは、余計なデータのコピーや型変換をまるごと省ける、それなりに "重い" 処理だと分かってきます。

こうして狙いを、正規化（GemmaRMSNorm）のまわりに絞り込んでいきました。

vLLM のソースを読んで、正体は GemmaRMSNorm だった

「正規化まわりが怪しい」と当たりをつけたので、vLLM のソースコードを読んでみます。Qwen3.5 が使っている正規化は GemmaRMSNorm というもので、モデルの各所——各ブロックの入力前、attention の後、そして最後の出力前——で繰り返し呼ばれていました。

そして、この正規化は内部で、次のような小さな処理の並びに分解されて実行されていました。

weight.float() + 1.0
x.to(float32)
pow → mean → rsqrt
multiply
to(orig_dtype)

（weight を float に直して +1、入力を float32 に変換、2 乗して平均をとり平方根の逆数を計算、掛け合わせて、最後に元の型に戻す……という手順です。）

これは、さきほど計測で見えていた「小さな処理が何度も並ぶ」パターンそのものです。正規化は層ごと・トークンごとに繰り返されるので、この分解のたびに、余計な中間データ（temporary tensor）の生成と、細かい kernel の呼び出しが積み重なります。ここを 1 つの kernel にまとめれば、無駄なメモリのやりとりと呼び出し回数を一気に減らせそうだ——これが有望な狙いどころでした。

自作した kernel — 分解された RMSNorm を 1 つにまとめる

やることはシンプルで、バラバラに分かれていた次の計算を、1 つの kernel にまとめます。

y = x * rsqrt(mean(x^2) + eps) * (weight + 1)

kernel の中身の方針はこうです。

• 1 行（= 1 トークン分の hidden vector）を 1 単位として処理する
• 平均や平方根の計算は、精度のため fp32（32bit 浮動小数点）で行う
• rsqrt(mean(x^2) + eps)（正規化のスケール）を計算する
• (weight + 1) を掛ける
• 最後に元の型に戻して書き出す

実装は Triton 版と CUDA C++ 版の両方を用意しました。vLLM への差し込みは、まず扱いやすい Triton 版で行い、あとから CUDA C++ 版も検証しています。組み込み方は、vLLM 本体のコードには手を入れず、起動時に元の正規化処理を自作版へこっそり差し替える形（いわゆる monkey patch）にしました。必要なときだけ有効にできるので、元の実装と切り替えて比較するのも簡単です。

まず単体で測る — mini benchmark

まずは、この kernel が単体でどれくらい速いかを測りました（torch native = 元の分解実行、triton/cuda fused = 自作版）。

shape	torch native	triton fused	cuda fused	best speedup
tokens=1, hidden=2048	37.888 us	13.840 us	7.168 us	5.286x
tokens=1, hidden=4096	39.776 us	13.536 us	6.928 us	5.741x
tokens=1, hidden=8192	38.976 us	13.312 us	7.168 us	5.437x
tokens=128, hidden=8192	51.200 us	13.456 us	8.192 us	6.250x

さきほど失敗した「単純な add/mul のまとめ」と違い、decode の tokens=1 でもはっきり勝てました。複数の処理を 1 kernel にまとめたことで、fp32 の中間データの生成やコピー・型変換が減り、平均の計算から最後の掛け算・書き出しまでをひと続きで処理できるためです。CUDA C++ 版がいちばん速いですが、本筋の backend 差し込みはまず Triton 版で確認し、CUDA C++ 版は追加の裏取りとして扱います。

vLLM に差し込んだ結果 — 本番の生成速度で確かめる

ここが本題です。単体で速いだけでは意味がなく、実際に vLLM へ差し込んで生成速度が上がるかどうかが勝負です。stream=True（1 トークンずつ返すモード）で max_tokens=128/512/2048 を計測しました（すべて最後まで生成しきった finish_reason=length の条件です）。

表の指標は、TTFT（最初のトークンが返るまでの時間）、TPOT（1 トークンあたりの生成時間）、ITL（トークンとトークンの間隔）、decode tokens/s（1 秒あたりに生成できるトークン数）です。

max tokens	variant	mean TTFT	mean TPOT	mean ITL p50	mean decode tokens/s
128	unpatched	34.092 ms	10.615 ms	10.618 ms	94.214
128	patched	30.048 ms	8.919 ms	8.897 ms	112.129
512	unpatched	34.173 ms	10.557 ms	10.550 ms	94.720
512	patched	29.531 ms	8.927 ms	8.918 ms	112.017
2048	unpatched	34.349 ms	10.625 ms	10.596 ms	94.121
2048	patched	30.378 ms	9.037 ms	9.025 ms	110.659

差分だけ取り出すと、こうなります（unpatched = 元の状態、patched = 自作 kernel を差し込んだ状態）。

max tokens	TPOT 削減	総レイテンシ削減	decode tokens/s
128	15.98%	15.88%	1.190x
512	15.44%	15.43%	1.183x
2048	14.95%	14.94%	1.176x

TPOT / ITL がどの条件でも一貫して速くなり、長めの 2048 トークンでも decode の速度改善が保たれました。これは「1 トークンごとの小さな処理の積み重ね」に効いた、という見立てを裏づけます。TTFT（最初の 1 トークン）も速くなっていますが、ここには prefill と最初の decode が混ざるので、単独では強くは主張しません。

なお、この約 15% は streaming（warmup=1, runs=3）での TPOT の数字です。より厳密な条件（warmup=3, runs=5）に揃えた非 streaming の測り直しでは、レイテンシで約 10%（tokens/s で約 1.11x）の改善でした。測り方によって数字の大きさは多少変わりますが、「decode が速くなる」という結論と改善の方向は一致しています。

Nsight で答え合わせ

狙ったところが本当に減ったのか、もう一度 Nsight で確認します。

family	baseline share	patched share
copy / cast	3.695%	1.443%
norm / reduce	2.708%	0.741%
elementwise	2.342%	0.484%

狙っていた copy/cast・norm/reduce・elementwise が、いずれもはっきり減りました。一方で cuBLAS GEMV の割合は 86.788% → 92.592% に上がっています。これは GEMV が遅くなったのではなく、周りの小さな処理が減って全体の時間が短くなった結果、もともといちばん割合の大きかった GEMV が、相対的にさらに大きな割合を占めるようになっただけです。「次に重いのは、また GEMV」という健全なプロファイルの変化で、streaming の速度改善とも符合します。

ここで、鋭い人は引っかかるかもしれません。「表で減ったのは全体の数％ぶんなのに、なぜ生成時間は 15% も速くなるの？」と。ポイントは、この表が測っているのが GPU が計算していた時間だけ、という点です。実際にはもう 1 つ、細かい処理を GPU に「これをやって」と 1 つずつ指示を出すたびに生じる、ちょっとした待ち時間があります。バラバラの処理を 1 つの kernel にまとめると、この指示出しの回数がごっそり減ります。今回の eager 実行はこの指示出しの待ちが積もりやすい状況だったので、そこを詰めた効果が、計算時間の削減以上に大きく効いた——というわけです。（裏を返すと、通常の実行ではこの指示出しがあらかじめまとめられているので、同じ効果は出にくくなります。後半の「限界」につながる話です。）

なぜ decode に効いたのか

RMSNorm は transformer の各ブロックで何度も実行されます。decode では 1 トークンずつ処理するので、レイヤー数のぶんだけ正規化と行列計算がトークンごとに発生します。Qwen3.5-2B のこの経路では、その正規化が小さな処理の並びに分解され、細かい kernel とコピー・型変換を生んでいました。

fusion（複数の処理を 1 つにまとめること）が効いた理由を整理します。

1. kernel の呼び出し回数が減る — バラバラだった複数の kernel を 1 つに。decode はトークンごとに繰り返すので、呼び出しのたびのオーバーヘッドが積み重なる。
2. メモリのやりとりが減る — fp32 の中間データや出力の型変換、途中の読み書きを減らし、GPU のメモリアクセスを削減する。
3. フレームワークによる分解を避けられる — 高水準の PyTorch の処理は、実行時に小さな kernel へ分解される。1 つは小さくても、レイヤー数 × トークン数で効いてくる。
4. decode という状況と相性がいい — batch が小さく 1 kernel の仕事量が少ないため、呼び出しオーバーヘッドが目立ちやすい。そこを削るので、1 トークンあたりの時間（TPOT / ITL）に直接効く。

単体の速さ ≠ backend の改善 — CUDA C++ 版での確かめ

mini benchmark では CUDA C++ 版がいちばん速いのでした。では、Triton 版で得た改善に、CUDA C++ 版でさらに上積みが出るのでしょうか？ 同じ条件で確かめました。

max tokens	variant	mean TPOT	mean decode tokens/s
128	unpatched	10.615 ms	94.214
128	Triton patched	8.919 ms	112.129
128	CUDA patched	8.879 ms	112.628
2048	unpatched	10.625 ms	94.121
2048	Triton patched	9.037 ms	110.659
2048	CUDA patched	9.039 ms	110.629

CUDA C++ 版も、元の状態に対しては TPOT 約 15% 改善を保ちました。ただし Triton 版からの上積みは 0.5% 未満で、ほぼ同じでした。単体では CUDA C++ がはっきり速くても、backend 全体では RMSNorm 以外の処理（特に GEMV）が時間の大半を占めるからです。

ここが大事な示唆です。単体 kernel の速さは、そのまま request 全体の改善としては出てこない。 だから本筋は Triton 版で成立させ、CUDA C++ 版はこの示唆を裏づける追加検証として扱いました。

限界と今後

この記事で主張できる範囲:

• Qwen3.5-2B / vLLM nightly / --enforce-eager / short prompt + decode-heavy
• stream=True/False の request-level benchmark
• Triton fused GemmaRMSNorm の monkey patch

まだ主張しないこと:

• すべての vLLM workload で速くなる
• CUDA Graph / torch.compile 有効時も同じ改善率が出る
• batch size > 1 / 長文 prefill / 他モデルでも同じ

補足です。--enforce-eager を外した通常の実行経路でも試しました。通常経路では CUDA Graph / torch.compile が小さな kernel の影響を吸収してしまうため、今回のような大きな TPOT 改善は出ませんでした。なので本記事は「通常経路を一般的に速くした話」ではなく、「eager 実行だと表に出てくる GemmaRMSNorm の分解を、自作 kernel で潰した話」として読んでください。

次に試すなら、この 3 点です。

• 通常経路で custom op を明示的に有効化（--compilation-config '{"custom_ops":["none","+gemma_rms_norm"]}'）し、fused path が確実に発火する条件で測り直す
• batch size > 1 / 長文 prefill での挙動を見る
• Qwen3.5-4B や他モデルでも再現するか確かめる

OSS に PR は出すのか — 意気揚々と、そっと閉じるまで

正直に書くと、この結果が出たときは、かなりテンションが上がっていました。「これは vLLM 本体に PR を出す価値があるのでは？」と、頭の中ではもう貢献者リストに自分の名前が載る妄想までしていたくらいです。ところが、いざ既存の Issue や PR を調べてみると——すでに より本質的な GemmaRMSNorm 最適化の PR（#42251） が提案されていました。しかも、私が手を入れた GemmaRMSNorm 1 つだけでなく、同じパターンで未 fusion のまま残っていた十数個の処理を、まとめて根本から直すものでした。要するに、今回のやり方は本家に出すという意味では、ほぼ出番がなかったわけです。振り上げた拳の行き場をなくし、この改善は OSS に出すのではなく、ローカルでの自己満足プロジェクトとして、そっとフォルダの奥に閉じることにしました。

PR を見送った理由を簡単にまとめると、こうです。

• すでにもっと包括的な PR がある — #42251 は同じ enforce_eager の問題を、GemmaRMSNorm を含む同種の処理十数個に対してまとめて直すので、今回の変更箇所も置き換えられる可能性が高い。
• 重複 PR と判断されやすい — コードが悪いのではなく「同じ方向の改善がすでに進んでいる」ため、"superseded by #42251"（#42251 に取って代わられた）となりがち。
• メンテナは包括的な方を選ぶ — 同じ場所を触る PR が複数あれば、より広くカバーする方が採られやすい。「速いから」だけでは通らない。

とはいえ、GemmaRMSNorm を調べ、実装し、ベンチマークで確かめた過程そのものは無駄ではありません。こうして記事にできましたし、今後の GPU 最適化にも十分活かせています。

まとめ

• kernel 単体 benchmark の速さだけでは、LLM backend での価値は決まりません。
• 先に実際の backend を計測し、cuBLAS が強い行列計算（GEMV）本体は避けて、その周りに残っていた GemmaRMSNorm の分解を狙って 1 つの kernel にまとめました。
• その Triton fused kernel で、Qwen3.5-2B の decode において TPOT 約 15% 改善 / decode tokens/s 約 1.18x を確認し、Nsight 上でも狙った処理（copy/cast/norm/reduce）が減ったことと符合しました。
• 計測から対象を選び、実装し、TPOT / ITL で backend への効果まで確認する——この一連のプロセスそのものが、本記事で示したかったことです。

コード・生データ・kernel 実装はリポジトリの decode_projection_fusion/ にあります。

https://github.com/CHIPMUNK-T0T/cuda-kernel-engineering