🛜

ZeroNIC の解剖メモ ~ バカな...はやすぎる... ~

に公開

はじめに

スタンフォード大学、コーネル大学、Enfabricaの共同研究による高性能ホストネットワーキング技術「ZeroNIC」について調べたメモ。

USENIX Symposium on Operating Systems Design and Implementation 2024で発表された論文「High-throughput and Flexible Host Networking for Accelerated Computing」をベースにしている。ZeroNICはAIやデータ分析などの帯域を大量に消費するアプリケーション向けのネットワーキング技術で、ハードウェアとソフトウェアを共同設計することで高性能と柔軟性の両立を実現している。

動画はこちら。
https://youtu.be/zdslmVqIsjc?list=PLbRoZ5Rrl5lcNznjOKrTgFmIsFRPA5da-

従来のネットワークスタックの課題

現代のデータセンターアプリケーションは、ネットワーク経由で大量のデータを扱う必要がある。例えばNVIDIA DGX-B200は各GPUあたり400Gbpsの帯域を持ち、システム全体で3.2Tbpsもの帯域を備えている。これだけのデータ転送量をさばくとなると、従来のネットワークスタックではいろいろと問題が出てくる。

スタック 特徴 問題
RDMA ・高スループット
・CPUを使わない
・NICハードウェアでプロトコル処理
・柔軟性がない
・無損失ファブリックが必要
・ハードウェア依存でプロトコル改善が困難
・ヘッドオブラインブロッキングなどの問題が解決しづらい
Linux TCP ・柔軟性がある
・様々な検証済みプロトコル
・新しいプロトコルも簡単に導入可能
・高スループットが出ない
・CPUがボトルネック
・カーネルとユーザー間のデータコピーが多すぎる

研究チームの実験によると、Linux TCPでは単一フローで50Gbpsが限界値。これはCPU使用率がほぼ100%になってしまうから。一方、受信側のゼロコピーを実装するとCPU使用率が大幅に下がり、100Gbpsまで向上する。つまり、データコピーの削減がCPU使用率低減の鍵というわけだ。

データパスとコントロールパスの分離アプローチ

ZeroNICの核心となるアイデアは、データパス(ネットワークとアプリケーションバッファ間のペイロード転送)とコントロールパス(ヘッダ処理とプロトコル制御)を物理的に分離すること。


https://www.usenix.org/system/files/osdi24-skiadopoulos.pdf

特徴をまとめるとこんな感じ。

コンポーネント 役割 特徴
NIC ・パケットのヘッダとペイロード分離
・ペイロードを直接アプリケーションバッファにDMA転送
・フロー状態の追跡と順序付け
ハードウェアで実装されるため高速かつ効率的
コントロールスタック ・NIC、プロトコル、アプリケーション間の調整
・メモリ管理とシグナリング
・キュー管理
設計の中核で、物理的に分離したデータパスとコントロールパスを論理的に結合
トランスポートプロトコル ・ヘッダのみを処理
・輻輳制御や信頼性の提供
・任意の環境で実行可能
・カーネル空間
・ユーザ空間
・アクセラレータなどで実行可能

面白いのは、この分離によって高性能と柔軟性の両方が実現できること。アプリケーションバッファがCPUにあってもGPUにあっても、トランスポートプロトコルがカーネルで動いてもユーザー空間で動いても、ZeroNICは対応できる。

NICハードウェアの設計

ZeroNICのNICハードウェアは、データパスとコントロールパスの分離を効率的に実現するため、独自のアーキテクチャを採用している。


https://www.usenix.org/system/files/osdi24-skiadopoulos.pdf

ハードウェアブロックの構成

この図は大きく分けて以下のコンポーネントで構成されている。

コンポーネント 役割
Transport Protocol Execution Environment トランスポートプロトコル(TCP等)が実行される環境
Application Buffers アプリケーションデータが格納される場所(この例ではGPU HBM)
NIC本体 パケット処理の中核部分
ネットワークポート 外部ネットワークとの接続点

NIC内部のコンポーネント詳細

NIC本体には様々なコンポーネントがあり、それぞれ重要な役割を担っている。

NICコンポーネント 機能
DMA ・ヘッダとペイロードの直接メモリアクセス処理
・CPUを介さない高速データ転送
TX Merge Unit ・送信時にヘッダとペイロードを結合
・完全なパケットの組み立て
RX Split Unit ・受信時にパケットをヘッダとペイロードに分割
・分離したデータの適切な転送先決定
MS List ・Memory Segment List
・アプリケーションの送受信バッファ情報を管理
・フローごとにリンクドリスト形式で保持
MR Table ・Memory Region Table
・アプリケーションが登録したメモリ領域を追跡
・IOMMUのアドレス変換情報を保持
Flow Table ・フロー状態の管理
・フローカーソル情報の保持
・パケットの処理決定に使用

NICハードウェアの動作で特に重要なのが、各テーブルとリストの連携だ。

  1. パケット到着時
    • パケットがネットワークポートから到着
    • RX Split Unitがパケットをヘッダとペイロードに分割
    • パケットのフロー情報(送信元・宛先IPアドレス/ポート、プロトコル)を使ってFlow Tableを検索
    • Flow Tableからフローカーソル(最後に処理された位置)とMS List情報を取得
    • MS Listから適切なメモリセグメント情報を特定
    • MR Tableを参照してアプリケーションメモリの物理アドレスを解決
    • ペイロードを直接アプリケーションバッファ(この例ではGPU HBM)にDMA
    • ヘッダをTransport Protocolに転送
  2. パケット送信時
    • Transport ProtocolがNICにヘッダを送信
    • MS ListとMR Tableを使って送信データの場所を特定
    • DMAdでアプリケーションバッファ(GPU HBM)からペイロードを読み取り
    • TX Merge Unitでヘッダとペイロードを結合
    • 完成したパケットをネットワークポートから送信

MS ListとMR Tableはゼロコピー転送の核心部分で、アプリケーションバッファの位置情報をNICに提供する。Flow Tableはパケットの順序管理や再送処理を効率化するためのデータ構造だ。これらのシンプルなデータ構造の組み合わせで、複雑なネットワーク処理を高速に実行できる。

送受信の流れ

ZeroNICの送受信フローを詳細に見ていこう。


https://www.usenix.org/system/files/osdi24-skiadopoulos.pdf

基本構成

この図は大きく分けて以下のコンポーネントで構成されている。

コンポーネント 役割
Application ユーザアプリケーション(PyTorch、Redisなど)
Provider Library アプリケーションとコントロールスタックの橋渡し役
Control Stack NIC、トランスポートプロトコル、アプリケーション間の調整役
Transport Protocol TCPなどのプロトコル処理を担当(コントロールパス)
NIC ネットワークインターフェース(データパスとコントロールパスの分離を実装)
Application Buffers CPUメモリ(DRAM)やGPUメモリ(HBM)などのデータバッファ

図中の数字は処理の順序を表していて、青い線は送信パス、赤い線は受信パスを示している。

送信パス(青い矢印)の流れ


https://www.usenix.org/system/files/osdi24-skiadopoulos.pdf

送信パスは以下の順序で処理される。

  1. 初期化
    アプリケーションはまずProvider Libraryを使って初期化。この段階で接続確立、メモリ登録などが行われる
  2. 送信リクエスト
    アプリケーションがsend(buff, len, ...)を呼び出す
  3. キュー管理
    Provider LibraryがControl Stackにリクエストを転送、Control StackはTransport Protocolに処理を委譲し、その後NICキューにエントリを追加
  4. プロトコル処理
    Transport Protocol(例えばTCP)がプロトコル固有の処理を実行
  5. ゼロコピーDMA
    NICがアプリケーションバッファから直接データをDMA転送(これがゼロコピーの核心部分)
  6. パケット形成
    NICがヘッダーとペイロードを結合して完全なパケットを形成
  7. 完了処理
    処理完了通知がNICからControl Stackを経由
  8. 通知転送
    Control StackからProvider Libraryへ完了通知
  9. アプリ通知
    最終的にアプリケーションに完了が通知される

受信パス(赤い矢印)の流れ


https://www.usenix.org/system/files/osdi24-skiadopoulos.pdf

受信パスは以下の順序で処理される。

  1. 初期化
    送信パスと同様、アプリケーションはまず初期化を行う
  2. 受信リクエスト
    アプリケーションがrecv(buff, len, ...)を呼び出す
  3. キュー管理
    Provider LibraryがControl Stack経由でNICに受信用のバッファ情報を伝える
  4. パケット処理
    NICがパケットを受信し、ヘッダとペイロードに分割
  5. ゼロコピーDMA
    ペイロードを直接アプリケーションバッファ(CPU DRAMやGPU HBM)にDMA転送
  6. ヘッダ転送
    ヘッダ情報をNICキューを通じてControl Stackに転送
  7. プロトコル処理
    Transport Protocolがヘッダのみでプロトコル処理(ACK生成など)
  8. 完了処理
    Control StackからProvider Libraryへ完了通知
  9. アプリ通知
    最終的にアプリケーションに完了が通知される

この設計の革新的な点

この図の中で特に革新的なのは以下の点だ。

革新点 説明
データパスとコントロールパスの分離 ・ヘッダ(コントロールパス)とペイロード(データパス)を物理的に分離
・プロトコル処理はヘッダだけで実行
・ペイロードは直接アプリケーションバッファに転送
ゼロコピー転送 ・カーネルバッファを経由せずに直接アプリケーションバッファへDMA
・CPUコピーオーバーヘッドの大幅削減
・CPU使用率の低減(100%→17%)
多様なエンドポイントサポート ・CPUメモリ(DRAM)
・GPUメモリ(HBM)
・その他のアクセラレータメモリにも対応可能
柔軟なプロトコル実装 ・Transport Protocolはカーネル空間、ユーザ空間、専用ハードウェアなど様々な環境で実行可能
・既存のTCPスタックをそのまま利用可能
・新しいプロトコルも容易に導入可能

この設計のおかげで、ZeroNICはRDMAのような高性能とTCP/IPのような柔軟性を両立できている。データパスの効率化によって、単一TCP流で100Gbps近いスループットを実現しながら、CPU使用率は17%程度に抑えられている。かなり画期的な設計だと思う。

性能評価

ZeroNICのプロトタイプ実装では、Xilinx Virtex UltraScale+ FPGAを使ったNICと、コントロールスタックにLinuxカーネルのTCPプロトコルを統合している。これをMellanox ConnectX-6 NICと比較した結果がかなり興味深い。

メトリック ZeroNIC MLX TCP
(送信側ゼロコピーあり)
MLX RoCE
単一TCP流のスループット 96.37 Gbps 50.63 Gbps 98.03 Gbps
CPU使用率(プロトコル処理) 17.20% 100% N/A
推定最大スループット 560.29 Gbps 50.63 Gbps N/A

実アプリケーションでの性能向上も素晴らしい。

アプリケーション 改善率 注目ポイント
NCCL (GPU通信) 2.66× ・GPU間直接通信をサポート
・TCPの堅牢性を維持
・アプリケーション修正不要
Redis (キーバリューストア) 3.71× ・CPUオーバーヘッド大幅削減
・RoCEに匹敵する性能
・既存アプリケーションをそのまま利用可能

特筆すべきは、ZeroNICが高いスループットを実現しながらも、TCPの堅牢性を維持していること。例えばパケットロスが1%発生しても、TCPの再送メカニズムのおかげで性能がほとんど低下しない。これはRoCEでは難しいことだ。

ZeroNICの革新的な特徴

ZeroNICのアプローチを整理すると、いくつかの革新的な特徴が浮かび上がる。

革新点 説明
データパスとコントロールパスの分離 ・ヘッダ(コントロールパス)とペイロード(データパス)を物理的に分離
・プロトコル処理はヘッダだけで実行
・ペイロードは直接アプリケーションバッファに転送
ゼロコピー転送 ・カーネルバッファを経由せずに直接アプリケーションバッファへDMA
・CPUコピーオーバーヘッドの大幅削減
・CPU使用率の低減(100%→17%)
多様なエンドポイントサポート ・CPUメモリ(DRAM)
・GPUメモリ(HBM)
・その他のアクセラレータメモリにも対応可能
柔軟なプロトコル実装 ・Transport Protocolはカーネル空間、ユーザ空間、専用ハードウェアなど様々な環境で実行可能
・既存のTCPスタックをそのまま利用可能
・新しいプロトコルも容易に導入可能
再送と順序変更への効率的対応 ・MS ListとFlow Tableを活用した効率的なパケット管理
・順序が入れ替わったパケットも適切な位置にDMA
・カーソル管理による一貫性確保

まとめ

ZeroNICは、RDMAのような高性能とTCP/IPのような柔軟性を兼ね備えた新世代のネットワークアーキテクチャだ。データパスとコントロールパスの物理的分離という斬新なアプローチにより、AI訓練や分散コンピューティングなどの帯域要求の厳しいアプリケーションにも対応できる。そして何より、既存のアプリケーションを修正することなく高性能化できるのが魅力的。次世代のデータセンターネットワーキングの有力候補として注目したい技術だ。

Discussion