初めに
結論
- 開発環境
対策: Depthformer と backbone を CUDA Graph 化する
- Depthformer ループ
- backbone

初めに

LFM2.5-Audio-1.5B-JP は、音声入力に対してテキストと音声を生成する日本語 speech-to-speech（s2s）会話モデルです。

前回までの記事で、このモデルに LoRA でずんだもん口調と声を追加学習しました。ただし実際に動かすとリアルタイムでは推論ができなかったので、これを高速化してリアルタイムに推論できるようにしてみました。

ayousanz.hatenadiary.jp

結論

1 生成ステップを分解すると、Depthformer の 8 codebook 逐次生成が全体の約 58% を占める単一最大のボトルネックでした。そこで Depthformer と backbone を CUDA Graph 化したところ、本番品質（top_k=4 サンプリング）で RTF 0.580 まで高速化できました。

開発環境

項目	バージョン / 構成
OS	Windows 11
GPU	RTX 4070 Ti SUPER
パッケージ管理	uv
Python	3.12
torch	2.11 + CUDA 12.8 wheel
liquid-audio	1.3.0
peft	0.19.1

Windows では Triton の公式ビルドがないため、torch.compile が実質使えません。

import os

os.environ.setdefault("NO_TORCH_COMPILE", "1")
os.environ.setdefault("NO_CUDA_GRAPH", "1")
os.environ.setdefault("PYTHONUTF8", "1")
os.environ.setdefault("PYTHONIOENCODING", "utf-8")

対策: Depthformer と backbone を CUDA Graph 化する

ボトルネックの大きい順に、Depthformer → backbone と 2 段階で Graph 化します。

Depthformer ループ

greedy 生成なら _sample_audio_frame（8 codebook を生成するループ全体）は静的な CUDA オペ列になるため、ループ丸ごとを単一の CUDA Graph に捕捉できます。ロード済みの model（LFM2AudioModel）に対して、次のように本物の _sample_audio_frame を「graph を replay する関数」へ差し替えます。

import torch

hidden = model.lfm.config.hidden_size
# replay のたびに値をコピーする固定入力。出力も固定アドレスに返る
static_emb = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)

# 1) warmup（capture 前に数回まわしてアロケータを安定させる）
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for _ in range(5):
        model._sample_audio_frame(static_emb, temperature=None, top_k=None)
torch.cuda.current_stream().wait_stream(s)
torch.cuda.synchronize()

# 2) 1 フレーム分（8 codebook 逐次）を丸ごと 1 つの graph に捕捉
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    static_out = model._sample_audio_frame(static_emb, temperature=None, top_k=None)

# 3) 「入力をコピー → replay → 出力を返す」で差し替える
def graphed_frame(embedding, *, temperature=None, top_k=None):
    static_emb.copy_(embedding)
    g.replay()
    return static_out.clone()

model._sample_audio_frame = graphed_frame

あとは通常どおり generate_interleaved を回すだけで、フレーム生成部が graph replay になります。end-to-end は次のとおりです（greedy。入力が同一ならフレーム数も一致するので、差はフレーム生成の速度だけです）。

条件	RTF	フレーム時間
eager-greedy	2.500	約 200ms/frame
graphed-greedy	1.051	約 84ms/frame

最大のボトルネック（58%）をほぼ消去でき、出力トークンは eager と完全一致しました。ただしこれ単独では RTF<1.0 を跨げません（1.051）。残る backbone + mimi が 80ms 予算を僅かに超えるためです。

backbone

backbone は decode のたびに KV キャッシュが伸びるため、そのままでは単一グラフを replay できません。ここで HF の StaticCache が LFM2 hybrid をネイティブ対応していることが効きました。固定アドレスのバッファに KV を書き込むので、prefill だけ eager で済ませ、seq=1 の decode step を 1 度だけ捕捉して各ステップ replay すれば正しく動きます。生成ループ内の self.lfm(...) 呼び出しを、次の関数に置き換えます。

import torch
from transformers.cache_utils import StaticCache

lfm = model.lfm
hidden = lfm.config.hidden_size
# 固定長バッファ。KV を固定アドレスに書くので decode step を graph 化できる
cache = StaticCache(config=lfm.config, max_cache_len=prefill_len + max_new_tokens + 8)

graph, static_out = None, None
static_in = torch.empty(1, 1, hidden, device="cuda", dtype=torch.bfloat16)

# warmup は KV を書き換えてしまうので、前後でスナップショット/復元する
def snapshot(c):
    return [{k: v.clone() for k, v in vars(l).items() if torch.is_tensor(v)} for l in c.layers]
def restore(c, snap):
    for l, d in zip(c.layers, snap):
        for k, t in d.items():
            getattr(l, k).copy_(t)

def lfm_decode_step(emb):  # emb: (1, 1, hidden)。prefill は eager で済ませた後の各 decode step
    global graph, static_out
    if graph is None:
        snap = snapshot(cache)
        s = torch.cuda.Stream(); s.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(s):
            for _ in range(3):
                static_in.copy_(emb)
                lfm(inputs_embeds=static_in, past_key_values=cache, use_cache=True)
        torch.cuda.current_stream().wait_stream(s); torch.cuda.synchronize(); restore(cache, snap)
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph):
            static_out = lfm(inputs_embeds=static_in, past_key_values=cache, use_cache=True)
        restore(cache, snap)
    static_in.copy_(emb)
    graph.replay()
    return static_out.last_hidden_state.clone()

ポイントは、position_ids と attention mask が cumulative_length（固定アドレスの GPU tensor）から毎ステップ再計算されるため、1 度捕捉した graph を replay するだけで各 decode step が正しくなることです。手動の buffer 実装も torch.compile（Triton）も要りません。backbone と Depthformer の両方を Graph 化した end-to-end は次のとおりです。

条件	RTF	lfm decode step
eager	2.292	28.30ms
+ depthformer graph のみ	1.156	28.23ms
+ backbone graph（both graph, greedy）	0.747	2.67ms

backbone graph で lfm decode step が 28.30→2.67ms（10.6 倍）になり、RTF < 1.0 を直接達成（0.747）しました。torch.compile（= Triton）を使わず生 CUDA Graph だけで成立するため、Windows でも動きます。

最後に本番品質（サンプリング）。本番のずんだ声生成は top_k=4 の multinomial サンプリングですが、CUDA Graph は乱数生成（multinomial）も捕捉するため、greedy と同じ graph がそのまま使えます。

条件	RTF	lfm step
eager	2.243	27.34ms
+ depthformer graph	1.114	27.75ms
+ both graph（greedy）	0.700	2.43ms
+ both graph（sampled top_k=4・本番品質）	0.580	1.21ms

yousanのメモ

LFM2.5-Audio-1.5B-JP の音声対話を CUDA Graph でリアルタイム化する

初めに

結論

開発環境

対策: Depthformer と backbone を CUDA Graph 化する

Depthformer ループ

backbone