AI Inference Engineering: 詳細な説明と実装例

1. 概要

チャットボットのAPIを呼び出して、レスポンスが画面に徐々に表示されるのを待つとき、あなたはLarge Language Model (LLM)のinferenceプロセスを目撃しています。この記事では、inferenceがtrainingとどう異なるのか、なぜそれが膨大なリソースを消費するのか、そしてエンジニアがそれをより高速かつ安価にするために使用する技術（KV cache、continuous batching、quantization、speculative decoding）について説明します。この記事を読み終わると、benchmarkのTTFT/TPOTという数字が何を意味するかを理解でき、実際のinferenceエンドポイントのパフォーマンスを測定するJavaコードが書けるようになります。

注: この記事ではtransformerの内部（attention、feed-forwardなど）には深く入り込みません。フォーカスはインフラエンジニアリングの観点です - requestが“inference”システムを通じてどのように“実行”されるかです。

2. Inferenceとは何か？

Trainingはデータから重み（weights）を学習するプロセスです。数週間にわたって数千のGPUを消費しますが、1回だけ実行されます。Inferenceは学習完了後のモデルで、promptを入力してoutputを取得します - これは毎回のリクエストで発生し、本番環境では毎日何百万ものリクエストが実行されます。したがって、trainingは1回の総コストとしては「高い」かもしれませんが、inferenceこそが長期の運用コストとユーザー体験（latency）を決定するものです。

重要な違いは、trainingではモデルが損失を計算するために全文を一度に見ることです。テキスト生成のinferenceでは、モデルは1つずつトークンを生成する必要があり、後続のトークンは前のトークンに依存します。これはautoregressive decodingと呼ばれます。

2.1. 2つのフェーズ：PrefillとDecode

inferenceリクエストは常に2つのフェーズを通過します:

• Prefill: モデルがすべての入力promptを一度に読み込み、すべての入力トークンのattentionを計算します。このフェーズはGPUを効率的に使用します（並列処理、compute-bound）。
• Decode: モデルが出力トークンを次々と生成し、各ステップで1つのトークンのみを計算します。このフェーズは順序的に実行されます（memory-bound - メモリ読み込み帯域幅のボトルネック、計算ではなく）。

Javaでシンプルなdecodeループの例（シミュレーション、実際のモデルではありません）で、その順序的な性質を見てみましょう:

public class AutoregressiveDecoder {

    public List<String> generate(List<String> prompt, int maxNewTokens) {
        List<String> context = new ArrayList<>(prompt);
        List<String> generated = new ArrayList<>();

        for (int step = 0; step < maxNewTokens; step++) {
            String nextToken = predictNextToken(context); // 1つのforward pass
            if (nextToken.equals("<EOS>")) break;

            context.add(nextToken);   // 新しいトークンが次のステップの入力となる
            generated.add(nextToken);
        }
        return generated;
    }

    private String predictNextToken(List<String> context) {
        // シミュレーション: 実際のモデルはtransformer全体をcontextに対して実行
        return "token_" + context.size();
    }
}

predictNextTokenを呼び出すたびに、attentionを計算するためにすべての既存のcontextを再度読み込む必要があります。最適化がなければ、コストは文の長さに応じて増加します。これがセクション4で紹介するKV cacheが生まれた動機です。

3. Inferenceのパフォーマンス測定: TTFTとTPOT

LLM servingをbenchmarkする際の2つの最も重要な指標:

• TTFT (Time To First Token): リクエスト送信から最初のトークンを受け取るまでの時間 - prefillフェーズのlatencyを反映します。
• TPOT (Time Per Output Token): decodeフェーズで各トークンを生成するのにかかる平均時間。

public record InferenceMetrics(Duration ttft, Duration avgTpot, int totalTokens) {

    public Duration totalLatency() {
        return ttft.plus(avgTpot.multipliedBy(totalTokens - 1));
    }
}

totalLatency()メソッドは実際の公式を示しています: 総時間 ≈ TTFT (1回) + TPOT × (残りのトークン数)。

200トークンの回答では、TPOTが単に10ms/tokenだけ遅いと、総時間に2秒が追加されます。これがservingのbenchmarkでTPOTがTTFTよりも詳しく検査される理由です。

実際の数値例（モデル7B、batch size 1、GPU A100）：
TTFT  ≈ 80ms
TPOT  ≈ 15ms/token
200トークン → 80ms + 15ms × 199 ≈ 3.07s

4. KV Cache: 既知の情報の再計算を避ける

cacheがなければ、各decodeステップはすべてのトークンシーケンスの最初からattentionを再計算する必要があります。複雑度は長さの二乗に比例して増加します。

KV cacheは処理済みトークンのkey/valueベクトルを保存し、次のステップで新しいトークンのみを計算して、古いcacheを再利用できます。

public class KvCache {

    private final Map<Integer, float[]> keyCache = new HashMap<>();
    private final Map<Integer, float[]> valueCache = new HashMap<>();

    public void store(int position, float[] key, float[] value) {
        keyCache.put(position, key);
        valueCache.put(position, value);
    }

    public List<float[]> allKeysUpTo(int position) {
        return IntStream.rangeClosed(0, position)
            .mapToObj(keyCache::get)
            .toList();
    }
}

storeは生成される新しいトークンごとに1回呼び出されます。allKeysUpToは次のdecodeステップが最初からforward passを再実行する必要なく、計算済みのすべてのkeyを取得できるようにします。

これが長いリクエスト（長いcontext）がVRAMを急速に消費する理由です。cacheはbatch_size × sequence_length × num_layers × hidden_sizeに比例します。

注: KV cacheは、context window（32K、128K token）がGPU RAMに膨大な負荷を与える主な理由です。計算ではなく。これはvLLMなどで使用されるpaged attentionなどの技術が生まれた理由でもあります。cache管理を仮想メモリ（page）のように行い、断片化を回避します。

5. Continuous batching: リクエストをグループ化してスループットを向上

GPUはバッチでの処理が個々のリクエストを処理するよりも効率的です。しかし、各リクエストの出力長は異なります。リクエストAは50トークンで完了し、リクエストBは300トークンまで続きます。

Static batching（バッチ全体が完了するまで待ってから新しいバッチを受け取る）はGPUを浪費します。最も遅いリクエストを待つ必要があるためです。

Continuous batchingはこれを解決します。バッチ内のリクエストが完了するとすぐに、その「スロット」を待機キューの新しいリクエストに置き換え、バッチ全体を待つ必要がありません。

public class ContinuousBatchScheduler {

    private final int batchSlots;
    private final Queue<InferenceRequest> waiting = new LinkedList<>();
    private final List<InferenceRequest> active;

    public ContinuousBatchScheduler(int batchSlots) {
        this.batchSlots = batchSlots;
        this.active = new ArrayList<>(Collections.nCopies(batchSlots, null));
    }

    public void onStep() {
        for (int slot = 0; slot < batchSlots; slot++) {
            InferenceRequest req = active.get(slot);
            if (req == null || req.isFinished()) {
                active.set(slot, waiting.poll()); // 待機キューから新しいリクエストを埋める
            } else {
                req.decodeNextToken();
            }
        }
    }
}

onStep()は各decodeステップで実行されます。スロットが空いているか、リクエストが完了した場合、waitingから次のリクエストを即座に受け取ります。

このおかげで、GPUは常にフルバッチで実行され、遅いリクエストを待つ「空き」スロットはありません。これはvLLM、TGI（Text Generation Inference）、およびTensorRT-LLMの中核メカニズムです。

6. Quantization: 精度とスピードのトレードオフ

元のモデルは通常FP16またはBF16（16ビット）で保存されます。Quantizationは保存ビット数（INT8、INT4）を減らしてVRAMを削減し、スループットを向上させますが、精度はわずかに低下します。

7Bパラメータモデル、ストレージ容量の比較:
- FP16: ~14 GB VRAM
- INT8: ~7 GB VRAM
- INT4: ~3.5 GB VRAM

vLLMなどのエンジンでサービス提供する場合、設定を宣言するだけです:

model: meta-llama/Llama-3-8B
quantization: awq
dtype: float16
gpu_memory_utilization: 0.9

quantization: awqはAWQ（Activation-aware Weight Quantization）を指定します。これはINT4技術で、最も重要な重みを優先的に保持することで、単純なquantizationよりも優れた精度を維持します。

注: 深いquantization（INT4以下）は、モデルの「幻覚」（hallucination）を増やしたり、推論能力を低下させる可能性があります。quantization後に必ずoutput品質をbenchmark（速度だけでなく）してください。特にcode generationや計算など、高い精度が必要なタスクの場合。

7. Speculative decoding: 事前予測でステップを削減

Decodeは順序的でスロー。各ステップで1つのトークンのみを生成するためです。

Speculative decodingは、小さく高速な“draft model”を使用して次のいくつかのトークンを予測し、その後、大規模な“target model”がそれらのトークンを複数回ではなく1回のforward passで検証（verify）するだけで済みます。

Draft modelが予測：「Hà」「Nội」「là」「thủ」「đô」
Target modelが1つのforward passで5つのトークンをすべて検証：
  - 4/5トークンが大規模モデルの分布と一致する場合 → 4トークンを保持、5番目のトークンをスキップ、そこからdecodeし直す。
  - 最初からミスの場合 → 通常のdecodeと同じ1つのforward passで済み、無駄がない。

この技術はoutput品質を低下させません（常にtarget modelで検証するため）。target modelで大規模なforward passを順序的に実行する回数を減らすだけです。draft modelが多くの予測を正確にする場合に最も効果的です。

8. 実践: LLM serving endpointからTTFT/TPOTを測定

vLLMサーバーがhttp://localhost:8000でOpenAI互換APIを実行していると仮定します。WebClientの依存関係を追加します:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-webflux</artifactId>
</dependency>

streaming responseを使用してTTFTを測定するクライアントを作成します:

@Service
public class LlmBenchmarkClient {

    private final WebClient webClient;

    public LlmBenchmarkClient(WebClient.Builder builder) {
        this.webClient = builder.baseUrl("http://localhost:8000").build();
    }

    public InferenceMetrics measure(String prompt) {
        Instant start = Instant.now();
        AtomicReference<Instant> firstTokenTime = new AtomicReference<>();
        AtomicInteger tokenCount = new AtomicInteger();

        webClient.post()
            .uri("/v1/completions")
            .bodyValue(Map.of("prompt", prompt, "stream", true, "max_tokens", 100))
            .retrieve()
            .bodyToFlux(String.class)
            .doOnNext(chunk -> {
                if (firstTokenTime.get() == null) {
                    firstTokenTime.set(Instant.now());
                }
                tokenCount.incrementAndGet();
            })
            .blockLast();

        Duration ttft = Duration.between(start, firstTokenTime.get());
        Duration total = Duration.between(start, Instant.now());
        Duration avgTpot = tokenCount.get() > 1
            ? total.minus(ttft).dividedBy(tokenCount.get() - 1)
            : Duration.ZERO;

        return new InferenceMetrics(ttft, avgTpot, tokenCount.get());
    }
}

doOnNextは新しいchunk（token）が到着するたびに呼び出されます。最初の呼び出しでfirstTokenTimeの時刻を記録し、これがTTFTです。

stream終了後（blockLast）、総時間からTTFTを引き、残りのトークン数で割ることで、平均TPOTを計算します。

8.1. JUnit 5およびAssertJでテスト

@SpringBootTest
class LlmBenchmarkClientTest {

    @Autowired
    LlmBenchmarkClient client;

    @Test
    void shouldMeetLatencySla() {
        InferenceMetrics metrics = client.measure("KV cacheを2文で説明してください。");

        assertThat(metrics.ttft()).isLessThan(Duration.ofMillis(500));
        assertThat(metrics.avgTpot()).isLessThan(Duration.ofMillis(50));
        assertThat(metrics.totalTokens()).isGreaterThan(0);
    }
}

このテストは回答の内容をチェックするのではなく、latencyのSLAのみをチェックします。これは“inference engineering”の正しいフォーカスです。システムが十分な速度で応答することを確認する（model品質の評価ではなく、それは別の範囲です）。

注: このテストでは、実際のLLMサーバーが実行されている必要があります（またはMockWebServerでWebClientをmockできます）。CI内では、このタイプのテストをintegration testとして分離し、通常のunit testと一緒に実行しないでください。外部インフラストラクチャに依存するためです。

9. よくあるエラーと自分でinference engineを構築しないべき時期

• Batching/KV cacheを最初から再度作成: ML インフラストラクチャの深い研究をしていない限り、vLLM、TGI、TensorRT-LLMが最適化したもので再度構築しないでください。彼らが既に持っているパフォーマンスに達するだけで数ヶ月かかります。
• スループットでのみbenchmark、tail latencyを忘れる: スループットが高くてもp99 latencyが悪いと、ユーザー体験は悪くなります。平均だけでなく、常にp50/p95/p99を測定してください。
• 品質をテストし直さないでQuantization: セクション6で述べたように、ビット削減には常にトレードオフがあります。速度を測定しながら品質測定をスキップするのは一般的な誤りです。
• 変動するトラフィックにstatic batchingを使用: リクエストの出力長に大きな差がある場合、static batchingはGPUを深刻に浪費します。最初からcontinuous batchingをサポートするエンジンを選択する必要があります。
• prefill-heavyとdecode-heavy workloadを分離しない: 本番環境の一部では、リソースを最適化するために、prefill処理ノードとdecode処理ノードを別々に分離します（disaggregated serving）。トラフィックに明確なパターンがある場合（例：長いpromptを持つRAG、短いoutput）、このアーキテクチャの検討は価値があります。

10. まとめ

Inference engineeringはlatency、throughput、およびGPUコストの最適化です。model trainingの問題ではありません。

prefill/decode、KV cache、continuous batching、quantization、speculative decodingを理解することで、benchmarkを読み解き、正しいserving engine（vLLM、TGIなど）を選択し、単に“実行速度が速いか遅いか”を見るだけでなく、正しい指標（TTFT、TPOT）を使用してシステムを測定できます。

自然な次のステップは、実際のconfigでquantizationまたはgpu_memory_utilizationを変更し、それらの指標がリアルなworkloadでどのように変化するかを確認することです。