トークンストリーミングサーバを自分で実装する

目標

SSEでトークンを流すサーバーを自分で作り、TTFTとトークン間のレイテンシを別々に測定し、クライアントの中断時に生成が実際に止まることを証明します。

なぜ重要なのか

総生成時間が同じ6秒でも、6秒後に一度に出てくるものと、200ms後から流れ出てくるものは、まったく別の製品です。そのため、LLMサービスでストリーミングは基本であり、その瞬間から、測定方法も変わります。応答の完了時刻1つではなく、最初のチャンクまでの時間と、チャンク間の間隔を、別々に測る必要があります。このラボで特に見落としやすいのが、ステップ7です。ユーザーが長い応答の途中で別の質問に移るパターンは非常によくありますが、キャンセルの処理がないと、それらのリクエストがバッチのスロットを占有し続けて、誰も見ていないトークンを作ります。同時スループットが理由なく下がり、GPUの請求書だけが増えます。静かで高くつく事故なので、1つのステップをまるごと割り当てました。

ステップ

/root/ls/gen.pyにgenerate(prompt, n)を作成してください。同じプロンプトとnなら、いつも同じトークン列を返す必要があります。/root/ls/gen1.txtと/root/ls/gen2.txtを同じ引数で作って、2つのファイルの内容が同じである必要があります。
/root/ls/server.pyを127.0.0.1:8170で起動してください。POST /generateは、{"prompt":"hello","max_tokens":20}を受け取り、{"text":"...","usage":{"prompt_tokens":<n>,"completion_tokens":20}}を返します。
GET /stream?prompt=hello&max_tokens=20は、Content-Type: text/event-streamで応答し、各トークンをdata: {"token":"..."}と空行で送ります。レスポンス本文に、data: で始まる行が20行以上ある必要があります。
ストリームの最後に、data: [DONE]を送ってください。
/root/ls/ttft.pyで最初のチャンクまでの時間を測り、/root/ls/ttft.txtにttft_ms=<정수>(プレースホルダーは整数です)を書いてください。0より大きく3000未満である必要があります。
/root/ls/itl.txtに、count=<n> p50_ms=<수> p99_ms=<수>(プレースホルダーは数値です)を書いてください。countは19以上である必要があります。
max_tokens=100でストリームを開き、5チャンクだけ受け取ったあとで切断してください。GET /_debug/generatedが、直前のストリームのリクエストで、サーバーが実際に作ったトークン数を返すようにして(累積の合計もいっしょに載せておくと便利です)、その値を/root/ls/cancel.txtに、requested=100 generated=<n>として書いてください。generatedは30以下である必要があります。
/root/ls/report.txtに、ttft_ms=<n>、itl_p50_ms=<수>、tokens=<n>、throughput_tps=<수>の4行を書いてください(プレースホルダーは数値です)。

参考

SSEのイベントは、data: <내용>(プレースホルダーは内容です)の1行のあとに、必ず空行が来ます。
フレームワークやプロキシのバッファリングをオフにしないと、ストリーミングが実際に流れません。
ストリームの途中のエラーは、ステータスコードで知らせられません。すでに200を送ったので、エラーもイベントとして送ります。
よくあるミス1: チャンクの間の空行を抜かして、クライアントが何も読めなくなることです。
よくあるミス2: キャンセルを検知せず、切れたリクエストのトークンを作り続けてしまうことです。

決定的なトークン生成器を作る

/root/ls/gen.pyにgenerate(prompt, n)を作成してください。同じプロンプトとnなら、いつも同じトークン列を返す必要があります。/root/ls/gen1.txtと/root/ls/gen2.txtを同じ引数で作って、2つのファイルの内容が同じである必要があります。

同じプロンプトに、いつも同じトークン列が出る必要があります。そうでなければ、採点も評価もできません。ハッシュをシードに使えば済みます。

非ストリーミングのエンドポイントを作る

/root/ls/server.pyを127.0.0.1:8170で起動してください。POST /generateは、{"prompt":"hello","max_tokens":20}を受け取り、{"text":"...","usage":{"prompt_tokens":<n>,"completion_tokens":20}}を返します。

まず、一度にまとめて返す形で作って、基準線を作ります。使用したトークン数もいっしょに返してください。

SSEでチャンクを送信する

GET /stream?prompt=hello&max_tokens=20は、Content-Type: text/event-streamで応答し、各トークンをdata: {"token":"..."}と空行で送ります。レスポンス本文に、data: で始まる行が20行以上ある必要があります。

コンテンツタイプと、イベントの区切りの規則が決まっています。イベントの間の空行を抜かすと、クライアントが認識できません。

終了シグナルを送る

ストリームの最後に、data: [DONE]を送ってください。

ストリームがいつ終わったかを知らせる、慣例的なマーカーがあります。

最初のトークンまでの時間を測る

/root/ls/ttft.pyで最初のチャンクまでの時間を測り、/root/ls/ttft.txtにttft_ms=<정수>(プレースホルダーは整数です)を書いてください。0より大きく3000未満である必要があります。

リクエストの開始から、最初のチャンクが届くまでの間です。全体の完了時刻と混同しないでください。

トークン間レイテンシの分布を出す

/root/ls/itl.txtに、count=<n> p50_ms=<수> p99_ms=<수>(プレースホルダーは数値です)を書いてください。countは19以上である必要があります。

チャンクごとに到着時刻を記録して、差を出します。平均だけでは足りません。

クライアントの中断時に生成を止める

max_tokens=100でストリームを開き、5チャンクだけ受け取ったあとで切断してください。GET /_debug/generatedが、直前のストリームのリクエストで、サーバーが実際に作ったトークン数を返すようにして(累積の合計もいっしょに載せておくと便利です)、その値を/root/ls/cancel.txtに、requested=100 generated=<n>として書いてください。generatedは30以下である必要があります。

切れたことを検知できなければ、止められません。サーバーが作ったトークン数を数えて、証明してください。

指標のレポートを作る

/root/ls/report.txtに、ttft_ms=<n>、itl_p50_ms=<수>、tokens=<n>、throughput_tps=<수>の4行を書いてください(プレースホルダーは数値です)。

前に測った値を、1つのファイルにまとめます。スループットは、生成したトークン数を総時間で割った値です。