DSparkがSGLangに統合:信頼度駆動による可変長検証

DSparkがSGLangに統合:信頼度駆動による可変長検証

Speculative Decodingの核心的なアイデアは、追加の計算によってデコードステップ数を削減することにある。しかし負荷が高まると、この計算コストは割に合わなくなる。バッチサイズB、1ステップあたりのドラフトトークン数がKの場合、ターゲットモデルは1ステップでB * K個のトークンを検証する必要がある。並行数がさらに増加すると、検証コストがデコードステップ削減による恩恵を上回る可能性がある。

DSparkはこの問題に両面から対処する。一方ではsemi-autoregressiveブロックドラフターを使用し、1回のドラフトフォワードで1トークンブロックを生成して高い受理率を維持する。もう一方では、ドラフトモデル自身の信頼度に基づいてリクエストごとに異なるverify lengthを動的に割り当て、受理される可能性が低い末尾トークンの検証を回避する。SGLangはDSparkをサポートし、Qwen3やDeepSeek-V4などのDenseモデルとSparseモデルの両方をカバーしている。

DSparkがSGLangで実現すること

今回の統合はsgl-project/sglang#30261に由来する。SGLangチームは論文の数値を厳密に再現しようとするのではなく、オープンなサービングエンジン上でDSparkの主要なメカニズムと性能曲線を再現することを目指している。具体的には、シングルユーザー時の速度向上、負荷増加時のverifyバジェット縮小、そしてこれらのスケジューリング戦略が実際のwall-clock性能にどう変換されるかを検証している。

DSparkのアルゴリズム側には主に3つのドラフト側コンポーネントが含まれる。

  • ブロックドラフター:Denseルート(Qwen3など)とSparseルート(DeepSeek-V4など)をサポート。1回のフォワードでgammaトークンのブロックを出力し、軽量なsequentialヘッド(MarkovまたはRNN)によって各ステップが前のトークンに条件依存するため、semi-autoregressiveとなる。
  • 信頼度ヘッド:各ドラフトトークンが検証を通過する確率を推定する。ブロック全体の確率積をblock survival probabilityとして利用できる。
  • Sequential Temperature Scaling(STS):信頼度を校正し、survival probabilityが実際のacceptance rateに近づくようにすることで、スケジューラーがバジェットを設定しやすくする。

SGLangはこれに加えて、サービング化に必要な機能を補完している。具体的には、Confidence scheduler、リクエストごとのragged verify、完全なCUDAグラフ、acceptance ceilingの可観測性、Additive SPSコストテーブル、Data-parallel attention、overlapスケジューラーの統合、そしてfused Triton kernelやshardedブロックドラフターのmatmulなどの性能最適化が含まれる。

混合トラフィック下:リクエストごとに同じverifyウィンドウを使うべきではない

均質なベンチマークはConfidence schedulerの価値を見えにくくしがちだ。実際の混合トラフィックでは、1つのバッチに高い予測可能性を持つ数学問題、オープンエンドの対話、詩の生成が混在する可能性がある。それぞれのacceptanceの難易度は異なるため、異なる長さのverifyウィンドウを与えるべきだ。

リクエストごとのverifyバジェット(左):cap-acceptモード下でのgsm8k、arena-hard、poetryにおけるverifyステップごとのceiling/window/delivered tokens。および3つのワークロードにおけるステップごとのverify長分布(右)。

上図は、gsm8k、arena-hard、poetryの3種類のワークロードにおけるcap-acceptモードでのverifyバジェットを示している。左側ではceiling、scheduler windowと実際のdelivered tokensを比較し、右側では各ステップのverify長の分布を示している。結果は、DSparkが異なるワークロードの受理特性に応じて異なるバジェットを割り当てられることを示している。受理率の高いタスクはより長いウィンドウを維持でき、受理率が低いまたは不安定なタスクは末尾を早めにカットして無効な検証を削減できる。

Overlapスケジューラー:スケジューリングオーバーヘッドをフォワードの後ろに隠す

DSparkの動的スケジューリングを適切に実装しないと、ドラフト生成とターゲット検証の間に余分なバブルが生じ、アルゴリズムの恩恵を相殺してしまう。SGLangはDSparkのspeculativeパスをoverlapスケジューラーに統合し、スケジューラーがDSpark固有の特別処理をほとんど必要とせず、スケジューリングのオーバーヘッドをフォワード実行後にできる限り隠蔽できるようにしている。

バッチサイズ1でのデコード:overlapスケジューラーオフ(上)では、run_batchイテレーション間およびドラフト生成とターゲット検証フェーズ間にバブルが発生する。オン(下)では、これらすべてが連続して実行される。

図では、バッチサイズ1のデコードシナリオにおいて、overlapスケジューラーをオフにするとrun_batchイテレーション間およびドラフト生成とターゲット検証の2フェーズ間に明らかなバブルが発生することが確認できる。オンにすると、関連フェーズが連続して実行され、追加のスケジューリングオーバーヘッドがほぼゼロに近づく。

SPSコストテーブル:スケジューラーがオンラインでverifyバジェットを選択できるようにする

Confidence schedulerは重要な問いに答える必要がある。各リクエストの現在のステップで何トークンを検証するのが最も効率的か、という問いだ。SGLangはオフラインプロファイリングで得られたAdditive SPSコストテーブルを使用してステップ時間をモデル化し、オンラインサービング時にスケジューラーがこのコストテーブルを参照して各リクエストにより適切なverifyバジェットを選択できるようにしている。

Additive SPSコストテーブルのフィット — raw step timeとフィット(a)およびスループット(b) — およびSPSで予測されたデコードステップ時間と実測値の比較(c)、DeepSeek-V4 on H200。

DeepSeek-V4 + H200の環境では、SPSコストテーブルはraw step timeをフィットし、デコードステップ時間を予測できる。チームは、現在のSPSとキャリブレーションはまだ第一版の近似であり、コンテキスト長がステップコストに与える影響を十分にモデル化できていない可能性があるため、スケジューラーの最適動作点にはまだ改善の余地があると指摘している。

総合性能:DSparkはMTPとnon-spec双方を上回る

DeepSeek-V4-Flash、H200、DP-attentionによる4並列の実験では、SGLangは3本の曲線を比較した。non-speculativeのフロアライン、MTP(EAGLE-styleのベースライン、1-1-2と3-1-4の設定で各バッチサイズの最良値を採用)、そしてDSparkだ。speculative configを除き、3グループの設定は統一されている。

H200 dp4における集計スループット対ユーザーあたりデコード速度、アームごとに1本の曲線:non-specフロア、MTP、DSpark。右上が優れている。各マーカーは3ラウンド平均のバッチサイズを表す。

図の横軸はユーザーあたりのデコード速度、縦軸は集計スループットであり、各点はバッチサイズ1から256の並行設定のいずれかに対応し、3ラウンドの平均値を示している。右上が優れた結果を意味する。結果は、DSparkが全並行数スキャン範囲においてより優れたスループットとレイテンシのトレードオフを提供し、MTPとnon-specフロアを明確に上回ることを示している。

3つのverifyモード:static、compact、cap-accept

SGLangにおけるDSparkのverifyモードは、後続のエンジニアリング設計を理解する上での核心的な概念だ。

  • static:各ステップで完全なドラフトブロックを検証する。full-blockベースラインとなる。
  • compact:スケジューラーが各リクエストのために選択したウィンドウのみを検証する。本番環境のパスとなる。
  • cap-accept:完全なブロックを検証するが、ウィンドウ内のトークンのみをコミットする。compactと同じ出力を持ちながら、full verifyが本来いくつのトークンを受理するかを露出させ、トリミングによって隠されたacceptance ceilingの計測に使用できる。

Ragged verify:CUDAグラフを真に小さくする。パディングして計算するのではなく

リクエストごとのverifyウィンドウと固定形状のCUDAグラフは本質的に相性が悪い。同じバッチ内で、あるリクエストが2トークンしか検証しない一方で、別のリクエストは6トークンを検証する可能性がある。すべてのリクエストを完全なブロック幅にパディングすると、トリミングした計算を再び戻すことになる。

固定形状デコードグラフはすべてのリクエストを完全なブロック幅にパディングする(N x W = 18セル、うち8セルがパディング)。Ragged compactグラフはスケジューリング済みトークンを1つのバッファにフロントパックし、合計値のみを最も近いキャプチャ済みティアに切り上げる(12セル、うち2セルがパディング)。両者ともパディング部分もフォワードを通すが、raggedはパディングされたセルを大幅に削減する。

SGLangのアプローチは、バッチのragged構造を保持し、トークン総数をグラフキーとして使用することだ。まず異なるリクエストの可変長トークンをコンパクトバッファにフロントパックし、次に合計数を最も近いキャプチャ済みティアに切り上げる。これにより、バジェットがトリミングされるとパック済みの合計がより小さいティアに下がり、DSparkがリプレイするのは実際のより安価なグラフとなる。つまり、masked full-width forwardで空回りするのではなく、attentionとMLPの行数が実際に削減される。

このパックバッファはcu_seqlensに類似したvarlen入力を使用するため、バックエンドの既存attentionカーネルを再利用できる。DeepSeek-V4を例に挙げると、モデル自身のsparse-MLAパスflash_mlaをそのまま使用でき、新たなカーネルを追加する必要がない。DP attentionの場合、各ランクが同一ティアを共有し、通常はいずれかのランクが必要とする最大ティアを採用し、同期して下降する。

動的スケジューリング対full-block:恩恵は主に高並行時に現れる

チームはまた、compact(ステップごとにSPS-argmaxバジェットを使用する動的トリミング)とno-trim(同じragged pathで実行されるstatic full-blockスケジュール)の初期A/Bテストも実施した。ここでのスケジューラーはまだ第一版のvanilla実装であり、メカニズムのエンドツーエンドの動作を証明することに重点を置いており、完全に調整された最終的な数値を示すものではない。

図6

傾向は明確だ。動的バジェットの恩恵は主に高バッチシナリオで現れる。バッチサイズ1では、ターゲット検証はトークン数が増加しても大幅に遅くなることはないため、トリミングによる節約は限られており、2つのパスは通常近似した結果となる。並行数が増加しスループットがプラトーに入り始めると、トリミングがステップを短縮でき、compactが徐々に差を広げる。受理率の低いワークロードでは末尾のトリミング可能なスペースが大きいため、恩恵が早く現れ、幅も大きくなる。

可観測性:トリミングが真の上限を隠してしまわないようにする

compactの副作用の1つは「観測を削除してしまう」ことだ。schedulerウィンドウ内の最初のいくつかの位置しか検証しないため、ブロック全体を完全に検証した場合に何トークンを受理できたかをシステムは知ることができない。このceilingがなければ、あるトリミングが合理的な節約なのか、本来コミットできたトークンを失っているのかを判断できない。

図7

cap-acceptはまさにこのために設計されている。完全なブロックを検証するが、ウィンドウ内のトークンのみをコミットするため、提出結果はcompactと一致しながら、full verifyのacceptance ceilingを露出させる。SGLangはさらに、リクエストごとの信頼度、キャリブレーション指標(ECEなど)などの可観測性データを提供し、オフライン分析を容易にしている。

companion runの追加実行を望まない本番環境向けに、SGLangはblock-accept estimatorも実装している。これは将来のステップにおけるターゲットトークンとそのlogprobsを活用して、トリミングされた反事実的な末尾を推定し、推定区間を提供する。この方法は、トリミングされたトラジェクトリとトリミングされていないトラジェクトリのアンカートークンが類似した性質を持つという仮定に基づいている。

結論

SGLangにおけるDSparkの価値は単なる「より高速なSpeculative Decoding」にとどまらない。信頼度駆動の可変長検証をサービス可能なシステム実装として実現したことにある。Semi-autoregressiveブロックドラフターが高い受理率を維持し、Confidence schedulerが各リクエストのverifyバジェットを制御し、ragged CUDAグラフがトリミングされた計算を確実に実行しないようにし、overlapスケジューラーとSPSコストテーブルがアルゴリズムの恩恵をエンドツーエンドのスループットとレイテンシの向上に変換する。

H200 + DeepSeek-V4-Flashのテストでは、DSparkはスループット/シングルユーザー速度の曲線でMTPとnon-specベースラインを上回った。高並行および低受理率のワークロードでは、動的トリミングの優位性が特に顕著だった。コストモデル、キャリブレーション、スケジューラー戦略がさらに調整されるにつれ、オープンなサービングエンジンにおけるDSparkにはさらなる性能向上の余地が残されている。