長いプロンプトで llama_decode がクラッシュする：GGML_ASSERT n_tokens_all <= n_batch の原因と修正

症状はいつも同じです。テストではすべて正常に動くのに、ユーザーが長い文書を貼り付けたり、2 ページの手紙をスキャンしたりするとアプリが消えます。Swift のエラーも、キャッチできる例外もありません。デバイスのログには次のようなメッセージが出ています。

GGML_ASSERT(n_tokens_all <= cparams.n_batch) failed

ggml のアサーションは abort() です。プロセス全体が終了するため、do/catch では防げません。llama.cpp のバージョンによって、大きすぎるバッチはこの abort を引き起こすか、llama_decode がエラーを返すかのどちらかになります。いずれにしてもプロンプトは処理されず、修正方法は同じです。

原因：2 つの異なる上限

llama.cpp のコンテキストには、混同しやすいサイズ関連のパラメータがいくつかあります。

llama.cpp は n_batch を n_ctx 以下に制限するため、コンテキストが小さいうちは問題に気づきません。問題が起きるのは、1 ページ分のテキストが収まるよう n_ctx を 4096 や 8192 に増やし、n_batch をデフォルトのままにしたときです。すると 3000 トークンのプロンプトはコンテキストには収まっても、バッチには収まりません。

よくある実装では、プロンプト全体を 1 回の呼び出しで llama.cpp に渡します。

var tokens = tokenize(prompt)        // 3000 tokens
tokens.withUnsafeMutableBufferPointer {
    _ = llama_decode(ctx, llama_batch_get_one($0.baseAddress, Int32($0.count)))   // 3000 > n_batch (2048)
}

llama_batch_get_one はすべてのトークンを 1 つのバッチにまとめます。llama_decode はバッチを n_batch と照合し、3000 は 2048 を超えています。短いプロンプトではこの行に到達しないため、このバグはテストをすり抜けてしまいます。

修正 1：n_batch を n_ctx に合わせる

どのみちプロンプトはコンテキストに収まる必要があるので、バッチをコンテキストと同じ大きさにします。

var cp = llama_context_default_params()
cp.n_ctx = 4096
// The whole prompt is decoded in one call. It has to fit into one batch,
// or llama.cpp aborts with GGML_ASSERT(n_tokens_all <= cparams.n_batch).
cp.n_batch = cp.n_ctx
cp.n_ubatch = min(cp.n_ctx, 512)
guard let ctx = llama_init_from_model(model, cp) else { throw LLMError.loadFailed }

n_ubatch は 512 以下のままにしてください。これは計算バッファのサイズを決めるため、スマートフォンでは大きなマイクロバッチは実際にメモリを消費します。一方、大きな論理バッチはメモリを消費しません。llama.cpp は引き続きプロンプトを 512 トークンずつ処理し、単にすべてのトークンを 1 回の呼び出しで受け付けるだけです。n_ubatch を n_batch より大きくしてはいけません。

ほぼすべてのアプリで使うべき修正はこれです。1 行で済み、生成コードを変更する必要もありません。

修正 2：分割してデコードする

自分の管理外のコードがコンテキストを生成している場合など、コンテキストのパラメータを変更できないときは、n_batch 以下のサイズに分けてプロンプトを渡します。自前のバッチを使う場合、位置（position）、シーケンス ID、どのトークンで logits を取得するかは自分で管理する必要があります。

/// Decodes `tokens` in chunks of at most `llama_n_batch(ctx)` tokens.
/// Only the very last token requests logits, which is all sampling needs.
func decodeInChunks(_ ctx: OpaquePointer, _ tokens: [llama_token]) -> Bool {
    let size = Int(llama_n_batch(ctx))
    var batch = llama_batch_init(Int32(size), 0, 1)
    defer { llama_batch_free(batch) }

    var start = 0
    while start < tokens.count {
        let end = min(start + size, tokens.count)
        batch.n_tokens = 0
        for i in start..<end {
            let k = Int(batch.n_tokens)
            batch.token[k] = tokens[i]
            batch.pos[k] = llama_pos(i)              // absolute position in the sequence
            batch.n_seq_id[k] = 1
            batch.seq_id[k]![0] = 0                  // everything belongs to sequence 0
            batch.logits[k] = i == tokens.count - 1 ? 1 : 0
            batch.n_tokens += 1
        }
        if llama_decode(ctx, batch) != 0 { return false }
        start = end
    }
    return true
}

押さえておくべき点：

• 位置は絶対値です。2 番目のチャンクは 0 ではなく位置 size から始まります。位置を振り直すと、エラーメッセージもなく attention が壊れます。
• 新しいプロンプトを始めるときは、先にメモリをクリアします。llama_memory_clear(llama_get_memory(ctx), true) を呼ばないと、位置が前の会話と衝突します。
• logits は最後のトークンだけ。すべてのトークンで logits を要求するとメモリと時間を無駄にします。llama_sampler_sample(sampler, ctx, -1) は最後の logits を読み取ります。
• バッチを解放します。llama_batch_init はメモリを確保するので、llama_batch_free と対にしてください。

プロンプトの処理後は、生成ループで引き続き llama_batch_get_one を使って 1 トークンずつ渡せます。位置を明示しないバッチは、コンテキストのメモリ内の最後の位置の続きから処理されます。

修正 3：デコード前に長さをチェックする

プロンプトがコンテキスト自体より長い場合は、どちらの修正も役に立ちません。n_ctx を超えたトークンには行き場がなく、さらに回答のための余地も必要です。まずトークン数を数え、わかりやすいエラーで失敗させます。

enum LLMError: Error {
    case promptTooLong(tokens: Int, limit: Int)
}

let tokens = tokenize(prompt)
let limit = Int(contextLength) - maxTokens        // leave room for the answer
guard tokens.count <= limit else {
    throw LLMError.promptTooLong(tokens: tokens.count, limit: limit)
}

これで長い入力は通常の Swift のエラーになります。UI では、テキストが長すぎることをユーザーに伝える、テキストを切り詰める、分割して要約する、といった対応ができます。トークン化はデコードに比べて高速なので、このチェックのコストは無視できる程度です。トークン化の関数は iPhone セットアップガイドにあります。

どれくらいの長さで多すぎるのか

目安として、英語では 1 トークンがおよそ 3〜4 文字に相当し、ドイツ語や日本語などではそれより少なくなります。手紙 1 ページはおよそ 400〜800 トークンで、そこに指示文と例が加わります。チャットテンプレートや few-shot の例はすぐに積み上がるので、推測ではなく実測すべき理由がここにもあります。GBNF のガイドで紹介している構造化出力用の文法はコンテキストを消費しませんが、システムプロンプト内のキーの説明は消費します。

n_ctx を増やすのもタダではありません。KV キャッシュはコンテキストに比例して大きくなります。1.6 GB のモデルを載せたスマートフォンでは、4096 トークンが妥当なデフォルト値で、1 ページ分のテキストに十分な余裕があります。

チェックリスト

1. コンテキストの作成時に n_batch = n_ctx とし、n_ubatch を 512 以下に設定する。
2. 新しいプロンプトの前に毎回メモリをクリアする。
3. トークン数を数え、tokens + maxTokens が n_ctx を超えるときは promptTooLong を投げる。
4. 意図的に長い入力を使い、実機でテストする。公式の xcframework にはシミュレータ用スライスがありません（シミュレータに関する注意を参照）。