推論系統/自動駕駛
FlashDriveがVLAの4段階推論を統合的に枝刈り、Alpamayo 1.5のシングルGPU遅延を151ミリ秒に短縮
FlashDriveは、動画のKVを再利用し、推論tokenのドラフト生成を並列化するとともに、フローマッチングで変化の小さいステップをキャッシュする。論文ではRTX PRO 6000上で4.7倍の高速化と、軌跡誤差がほぼ変わらないことを報告しているが、現時点ではシミュレーション環境と単一のモデルスタックで得られた結果にとどまる。

自動運転向けVision-Language-Actionモデルの遅延は、単一の中核処理に集中しているわけではなく、視覚エンコーディング、言語モデルのprefill、推論tokenのデコード、フローマッチングによるアクション生成に分散している。そこで[FlashDrive](https://arxiv.org/abs/2608.12932)は、単にattention kernelを置き換えるのではなく、アルゴリズムとシステムを協調設計し、4つの各段階からそれぞれ異なる形の重複計算を排除する。
最初の最適化では、連続するマルチカメラ入力の時間的な重複を利用する。4フレームのスライディングウィンドウでは、通常、更新のたびに3フレームが変化しないため、システムは新しいフレームだけをエンコードする。さらにpre-RoPE keyを保存し、位置が移動した際に新たなrotary embeddingを適用する。カスタムマスクと組み合わせることで、視覚処理とprefillの有効シーケンス長を約75%削減する。古いKVは異なるコンテキストで計算されているため、そのまま再利用すると分布シフトが生じる。そこで著者らはVLMを凍結し、誤差の影響を受けやすいaction expertのみをfine-tuningして、軌跡精度を回復させている。
推論段階では、ブロック拡散モデルDFlashを使い、低エントロピーで定型的な運転推論tokenのシーケンス全体を一度にドラフト生成し、その後、元のモデルで検証する。アクション側では、フローマッチングのvelocity fieldがU字型に変化することが観察された。最初と最後のステップでは変化が大きい一方、中間部分のcosine similarityは0.99を超える。このため、端点付近のステップだけを再計算し、中間部分のvelocityをキャッシュする。最後に、W4A8 ParoQuant、CUDA Graph、およびQKVとMLPのkernel fusionを組み合わせる。
論文の完全版では、RTX PRO 6000上でAlpamayo 1.5-10Bの遅延を717ミリ秒から151ミリ秒へ短縮し、処理レートを1.4 Hzから6.6 Hzへ向上させたと報告している。[email protected]の変化はわずか0.08メートルだった。これより前の[プロジェクトページ](https://z-lab.ai/projects/flashdrive/)では159ミリ秒と報告されており、公開値がバージョンやテスト設定に応じて現在も調整されていることを示している。
より重要な制約として、コードとweightはまだ公開されておらず、closed-loopの結果もシミュレーション環境で得られたものに限られる。また、6.6 Hzという数値だけでは、実車制御、安全上の冗長性、worst-case latencyの要件を満たすことを証明できない。エンジニアリングチームが次に注目すべき点は、再現可能な実装、長いシーケンスにおけるキャッシュのドリフト、そしてJetson Thor上での消費電力とtail latencyである。