nixl

EFA の同じノード内の GPU 間通信: libfabric の shm と SAR はじめに AWS の GPU インスタンスで EFA を使う通信ライブラリの多くは、libfabric という通信のライブラリを通して EFA を操作します。たとえば NCCL は、aws-ofi-nccl というプラグインを通して libfabric を使い、libfabric が rdma-core とドライバを通して EFA を動かします。この層の分担は「MoE の all-to-all を支える層: DeepEP、NCCL GIN、aws-ofi-nccl、EFA」で図にしています。この記事で扱う NIXL の LIBFABRIC バックエンドも、同じように libfabric を通して EFA を使います。 通信のライブラリ libfabric は、通信の方式をプロバイダという部品で切り替えます。AWS のネットワークカードである EFA を使うのが efa プロバイダで、同じノードの中で共有メモリを使うのが shm プロバイダです。efa プロバイダを使うと、相手が別のノードにいるときは EFA の装置を通ってデータが届きます。ところが、この記事で調べた libfabric 2.4.0 の efa プロバイダは、相手が同じノードにいるとき、既定の設定 (FI_EFA_ENABLE_SHM_TRANSFER=1) では EFA の装置を使わず、shm プロバイダにデータを任せます。shm プロバイダは、1 つの登録済みの GPU のバッファを送るときなど、条件がそろうと CUDA IPC を選びます。CUDA IPC は、別のプロセスの GPU のメモリを、自分のプロセスから開いてコピーできるようにする CUDA の仕組みです。このとき、受け取る側が送る側の GPU のメモリを開いてコピーします。CUDA IPC を選んだあとで、送る側がその準備に失敗すると、SAR という、ホストの共有メモリを経由して小分けに運ぶ方法に切り替えます。この記事では、この同じノードの中の経路を図で追います。 既定では、同じノードの相手へのデータは EFA の装置を通りません 図の左は、相手が別のノードにいる場合です。データは送る側の EFA の装置からネットワークを通り、相手の EFA の装置へ届きます。図の右は、相手が同じノードにいる場合です。efa プロバイダは相手が同じノードにいると判断すると、既定の FI_EFA_ENABLE_SHM_TRANSFER=1 では shm プロバイダにデータを任せます。FI_EFA_ENABLE_SHM_TRANSFER=0 のときは、同じノードの相手にも EFA の装置を使います。RDMA (Remote Direct Memory Access) の書き込みは、通信の装置が、相手のメモリへ直接データを書き込む操作です。測った構成では、この転送が装置の RDMA の書き込みとして出たことを、「直し方 1」の節で確かめます。 ここから、この記事で使う 3 つの言葉を、図で 1 つずつ説明します。 shm プロバイダは、ホストのメモリの領域で命令と応答を受け渡します shm プロバイダでは、各プロセスがホストのメモリに自分の領域を作ります。送る側は、相手が作った領域を shm_open で開いて mmap で自分のアドレス空間に対応付け、その領域にある命令の列 (command queue) に命令を書きます。受け取る側は、自分の領域に書かれた命令を読んで処理し、処理が終わったことを、送る側の領域にある応答の列 (response queue) で返します。受け取る側の領域には、あとで見る SAR で使う中継用の領域も入っています。データの中身がどう運ばれるかは、次の CUDA IPC と SAR で違います。 CUDA IPC では、受け取る側が送る側の GPU のバッファから 1 回でコピーします CUDA IPC (Inter-Process Communication) は、別のプロセスの GPU のメモリを使えるようにする CUDA の仕組みです。送る側は、自分の GPU のバッファから IPC ハンドルという識別子を作り、shm の命令に入れて渡します。この記事で扱う、GPU のバッファへの送信では、受け取る側がそのハンドルを開き、送る側の GPU のバッファから自分の GPU のバッファへ、1 回でコピーします。コピーをするのは、受け取る側のプロセスです。ハンドルを作る前に、送る側は cuMemGetAddressRange でバッファの場所を調べます。この呼び出しは、送る側のスレッドで有効な CUDA のコンテキスト (GPU ごとに作られ、その GPU で確保したメモリを管理するもの) の中でアドレスを探すので、そのコンテキストの中でバッファが見つからないとハンドルを作れません。この記事で見る失敗は、ここで起きます。手順の細かいところは「CUDA IPC による受け渡し」の節で見ます。 SAR では、ホストの共有メモリを中継して、1 かたまりにつき 2 回コピーします SAR (segmentation and reassembly) は、shm がデータを中継用の領域で小分けに運ぶ方法です。データの大きさや指定によって最初から選ばれることもあり、CUDA IPC を選んだあとでその準備に失敗したときの切り替え先にもなります。この記事で問題になるのは、後者の切り替えです。SAR では、データを決まった大きさ のかたまりに分け、送る側が自分の GPU のバッファから、受け取る側の共有メモリにある中継用の領域へコピーします。続いて受け取る側が、そこから自分の GPU のバッファへコピーします。中継用の領域は使い回し、データが全部渡るまでこれを繰り返します。CUDA IPC なら 1 回で済むコピーが、かたまりごとに 2 回になり、そのたびにホストのメモリを経由します。手順の細かいところは「SAR」の節で見ます。 同じノードの中のこの経路で、実際に何が起きていたのかを、次の節で問題として定めます。 調べたケース: 同じノード内の LIBFABRIC 転送 この記事では、SGLang で disaggregated inference (分離推論) をするときに、NIXL の LIBFABRIC バックエンドを使って、同じノードにいる相手へ KV cache を送るケースを調べました。SGLang の分離推論では、入力を読む prefill と、出力のトークンを 1 つずつ作る decode を、別々のプロセスに分けて動かします。prefill は、入力の各トークンについて計算した Attention の Key と Value を KV cache として保存し、decode へ送ります。decode はそれを使って、入力を計算し直さずに続きのトークンを作ります。その転送を受け持つのが、推論のプロセスのあいだでデータを転送するライブラリ NIXL です。この記事の構成では、prefill と decode は、それぞれモデルの計算を 2 枚の GPU に分けて受け持ちます。この分け方をテンソル並列 (TP) と呼び、GPU ごとに 1 つずつあるプロセスの番号を TP rank と呼びます。prefill と decode は別々のノードに置くのが基本の構成なので、同じノードに置くケースは、あまり想定されていないかもしれません。今回は、このケースで何が起きるかを確かめます。 ! SGLang の分離構成で NIXL の LIBFABRIC バックエンドを使い、prefill と decode が同じノードにあると、KV cache は EFA の装置ではなく libfabric の shm の経路を通ります。測った構成の既定の設定 (NIXL 1.1.0) では、この経路の転送が、shm を切ったときより大幅に遅くなりました。この記事は、その経路で何が起きているかを追います。 NIXL は、転送に使うバックエンドを選べます。この記事で扱うのは LIBFABRIC のバックエンドです。EFA のあるインスタンスで efa プロバイダを使う場合、送り先が同じノードにあれば、KV cache は上の shm の経路を通ります。NIXL には UCX のバックエンドもあります。UCX は libfabric を使わず、同じノードの相手には、UCX 自身の経路 (cuda_ipc や共有メモリなど) から選んで使います。そのため、この記事の話は、NIXL で LIBFABRIC のバックエンドを選んだときのものです。UCX のバックエンドでも、転送スレッドが選ぶ GPU と KV cache のある GPU が食い違う問題が起きるかどうかは、この記事では調べていません。 測った構成では、shm を使う既定の設定で、SGLang 自身が記録する KV 転送時間が、shm を使わない設定より大幅に長くなりました。値は最後の「4 つの組み合わせの値」の節に並べます。page size は、SGLang が KV cache を管理するページ 1 つに入るトークンの数で、この記事の数値はすべて page size 1 で測ったものです。この記事が答える問いは、次の 2 つです。 shm の経路のどこで、何が起きているのか それを避ける方法は何か 測った構成の shm の経路では、prefill の TP rank 1 (送る側) で、KV cache の書き込みを発行する SGLang の転送スレッドが、データのある 1 番の GPU ではなく 0 番の GPU を向いたまま、書き込みを発行していたと考えられます。CUDA は、いま使う GPU (current device) をスレッドごとに覚えていて、転送スレッドではこれを設定していません。そのため CUDA IPC の準備に失敗し、SAR に切り替わります。転送スレッドで GPU を指定すると、この失敗は起きなくなりました。この記事では、これを「問題」と呼びます。 この問題を避ける方法は、2 通りあります。 直し方 何を変えるか 1. shm を使わない設定 同じノードの相手にも、EFA の装置で RDMA の書き込みをする。shm の経路を使わない 2. NIXL 1.4.0 shm の経路のまま、書き込みの前に、メモリを登録したときのコンテキストにスレッドを合わせる。SAR に切り替わらない 次の表は、転送に関わる層を、上から呼び出す順に並べたものです。右の列が、この記事で見るところです。 層 受け持ち この記事で見るところ SGLang KV cache を持ち、転送を頼む 転送スレッドが GPU を選ばないこと NIXL の LIBFABRIC バックエンド 転送を libfabric の呼び出しに変える 送る前の CUDA のコンテキスト libfabric の efa 相手が同じノードなら shm に任せる shm を使う設定 libfabric の shm 共有メモリで渡す CUDA IPC と SAR 記事の中の「0 番」「1 番」は、各プロセスから見た GPU の番号を指します。prefill のプロセスには CUDA_VISIBLE_DEVICES=0,1、decode のプロセスには CUDA_VISIBLE_DEVICES=2,3 を設定していたので、たとえば decode の「0 番」はノードの 3 枚目の GPU です。この記事の数値は、次の条件で測ったものです。 調べた条件 項目 値 インスタンス g7.48xlarge 1 台。Kubernetes の 1 つの Pod に prefill と decode を入れ、それぞれテンソル並列の TP 2 で 2 枚の GPU を使う。Pod に割り当てた EFA の装置は 1 台で、prefill と decode が一緒に使う モデル DeepSeek-V2-Lite-Chat。revision は 85864749cd611b4353ce1decdb286193298f64c7 SGLang 0.5.12.post1 NIXL 1.1.0 と 1.4.0 libfabric 2.4.0 ルーター sglang-router 0.3.2、--pd-disaggregation --policy round_robin --disable-retries --request-timeout-secs 330 prefill と decode の共通の引数 --tp-size 2 --dtype bfloat16 --kv-cache-dtype bf16 --context-length 32768 --chunked-prefill-size 8192 --mem-fraction-static 0.75 --page-size 1 --enable-metrics --stream-interval 1 --disaggregation-transfer-backend nixl。prefix cache は有効のまま prefill だけの引数 --disaggregation-mode prefill --disaggregation-bootstrap-port 8998 decode だけの引数 --disaggregation-mode decode 環境変数 SGLANG_DISAGGREGATION_NIXL_BACKEND=LIBFABRIC、FI_PROVIDER=efa、FI_EFA_USE_DEVICE_RDMA=1、FI_LOG_LEVEL=info。shm を使う構成では FI_EFA_ENABLE_SHM_TRANSFER を設定せず既定の 1 のまま、shm を使わない構成では 0 要求 入力 8,192 トークン、出力 256 トークン。問題の警告の行は、起動したばかりのエンジンに要求を 1 つ送って数えた。KV 転送時間は、平均して 1 秒に 1 件の割合で 20 秒送って、構成ごとに 1 回測った。TTFT は、同じ負荷を、要求の届く間隔などをばらつかせる乱数の始まりの値 seed を 3 通りに変えて流し、合わせた 測ったもの SGLang が --enable-metrics で記録する段階ごとの時間。sglang:per_stage_req_latency_seconds の stage="prefill_transfer_kv_cache"。libfabric のログの警告の行。unable to use IPC for msg, fallback to using SAR と Failed to perform cuMemGetAddressRange 問題の警告は、手を加えていない SGLang の prefill と decode のサーバーを同じノードで起動し、ルーターを PD モードで立てて、長い入力を送るだけで出ます。KV 転送時間は、同じ長さの要求を、届く間隔をばらつかせながら平均して決まった割合で送り、測りました。 同じ構成を起動する手順 MODEL_PATH に、モデルを置いた場所か Hugging Face の名前を入れます。これは、上の条件の起動引数を並べたものです。このブロックは、上の条件と同じ版 (SGLang 0.5.12.post1、NIXL 1.1.0、libfabric 2.4.0) が入った、GPU 4 枚と EFA の装置 1 つを割り当てた 1 つの Pod の中で、そのまま流すと下に書いた数が出ることを、2026-10-20 に確かめています。最後の行で、起動したサーバーを止めます。 export MODEL_PATH=deepseek-ai/DeepSeek-V2-Lite-Chat export SGLANG_DISAGGREGATION_NIXL_BACKEND=LIBFABRIC export FI_PROVIDER=efa export FI_EFA_USE_DEVICE_RDMA=1 export FI_LOG_LEVEL=info COMMON="--model-path $MODEL_PATH --revision 85864749cd611b4353ce1decdb286193298f64c7 --host 127.0.0.1 --tp-size 2 --dtype bfloat16 --kv-cache-dtype bf16 --context-length 32768 --chunked-prefill-size 8192 --mem-fraction-static 0.75 --page-size 1 --enable-metrics --stream-interval 1 --disaggregation-transfer-backend nixl" CUDA_VISIBLE_DEVICES=0,1 python3 -m sglang.launch_server $COMMON --port 30000 --disaggregation-mode prefill --disaggregation-bootstrap-port 8998 > prefill.log 2>&1 & CUDA_VISIBLE_DEVICES=2,3 python3 -m sglang.launch_server $COMMON --port 30001 --disaggregation-mode decode > decode.log 2>&1 & until curl -sf 127.0.0.1 > /dev/null && curl -sf 127.0.0.1 > /dev/null; do sleep 10; done python3 -m sglang_router.launch_router --pd-disaggregation --prefill 127.0.0.1 8998 --decode 127.0.0.1 --host 127.0.0.1 --port 8000 --policy round_robin --disable-retries --request-timeout-secs 330 > router.log 2>&1 & until curl -sf 127.0.0.1 > /dev/null; do sleep 2; done python3 -c 'import json, random, urllib.request; ids = [random.randrange(1000, 30000) for _ in range(8192)]; req = urllib.request.Request("127.0.0.1", json.dumps({"input_ids": ids, "sampling_params": {"max_new_tokens": 256, "ignore_eos": True}}).encode(), {"content-type": "application/json"}); print(json.loads(urllib.request.urlopen(req, timeout=600).read())["meta_info"]["completion_tokens"])' grep -c "fallback to using SAR" prefill.log grep -c "cuMemGetAddressRange" prefill.log kill $(jobs -p) このブロックを流すと、生成されたトークンの数として 256 が出て、そのあとの 2 つの grep -c が、SAR への切り替えの警告と cuMemGetAddressRange の失敗の警告の行の数を出します。どちらも 0 より大きい数が出ます。shm を使わない構成を試すときは、サーバーを起動する前に export FI_EFA_ENABLE_SHM_TRANSFER=0 を足します。足して流すと、256 のあとに 0 と 0 が出ました。 以下では、次の「全体像」の節で層と経路の図を示したあと、efa から shm への受け渡し、CUDA IPC と SAR、問題 (送る側のスレッドの GPU)、直し方 1 の shm を使わない設定、直し方 2 の NIXL 1.4.0 の順に追い、最後に 4 つの組み合わせの値を並べます。 全体像 同じノードの相手へのデータは、既定では EFA の装置ではなく shm プロバイダを通ります SGLang は TP の rank ごとに 1 つのプロセスを立て、各プロセスは自分の GPU の上に KV cache を持ちます。KV cache の転送は NIXL が受け持ち、相手が同じノードにいるか別のノードにいるかで、libfabric の中の経路が分かれます。 NIXL は書き込みに、書き込みが終わったときに受け取る側へ一緒に届く小さな値を付けて発行します。この値を付けることを示す印が、libfabric の FI_REMOTE_CQ_DATA です。発行に使う関数は、1.1.0 では fi_writedata、1.4.0 では fi_writemsg です。この値の有無は、あとで見る shm の経路の選び方に関わります。 efa から shm への受け渡し 書き込みに印を付けるのは NIXL で、その印を見て経路を決めるのは libfabric です SGLang は KV cache の送信を NIXL に頼み、NIXL の LIBFABRIC バックエンドは、libfabric の書き込みの関数を FI_REMOTE_CQ_DATA の印付きで呼びます。その先の経路は libfabric が決めます。efa プロバイダは、相手が同じノードなら shm プロバイダに任せます。shm プロバイダは、印が付いているので送る側だけで書き終える CMA の経路を使わず、この記事の書き込みでは CUDA IPC か SAR で運びます。 efa プロバイダは、同じノードの相手との通信のために、shm プロバイダを中に抱えています。shm を使うかどうかは設定 enable_shm_transfer で決まり、既定値は 1 です。環境変数 FI_EFA_ENABLE_SHM_TRANSFER=0 でこれを 0 にすると、shm の経路そのものが作られません。この設定を使う直し方は、「直し方 1: shm を使わない設定」の節で見ます。 CUDA IPC による受け渡し shm の CUDA IPC では、データをコピーするのは受け取る側のプロセスです 別のプロセスの GPU のメモリは、そのアドレスを渡されただけでは使えません。CUDA IPC は、送る側が自分の GPU のメモリから IPC ハンドルという識別子を作り、受け取る側がそのハンドルを開いて、相手のメモリを自分のプロセスから使えるようにする仕組みです。 shm には、プロセス間のメモリコピー (CMA) を使って送る側だけで書き終える経路もあります。ただし、FI_REMOTE_CQ_DATA の付いた書き込みはこの経路を使いません。どの方法で渡すかは smr_select_proto が決め、送るバッファが 1 つで、登録済みの GPU のメモリにあり、相手も IPC を使えるときは、CUDA IPC を選びます。ここで登録とは、転送に使うメモリの領域を、通信のライブラリにあらかじめ知らせておく処理のことです。 細かい手順の前に、全体の流れを書きます。データは、送る側 (prefill) の GPU にある KV cache から、受け取る側 (decode) の GPU のバッファへ流れます。コピーを実行するのは、受け取る側のプロセスです。そのための通り道は、2 つの処理でできます。送る側が自分の GPU のバッファの IPC ハンドルを作って渡す処理と、受け取る側がそのハンドルを開いて送る側のバッファを読めるようにする処理です。ハンドルや完了の知らせは、ホストの共有メモリにある命令の列と応答の列でやり取りします。 図の番号の順に、1 回の書き込みで何が起きるかを見ます。libfabric では、完了を確かめる関数を呼んだときに、たまっている処理をその場で進めます。これを progress と呼びます。 送る側の転送スレッドが、cuMemGetAddressRange で、転送するポインタを含む確保領域の先頭と長さを調べます。この呼び出しは、そのスレッドでいま有効な CUDA のコンテキストの中でアドレスを探します。 その先頭について cudaIpcGetMemHandle で IPC ハンドルを作ります。ハンドルは確保領域の全体を指すので、転送するポインタの先頭からのずれ (オフセット) を一緒に記録します。 送る側は、ハンドル・オフセット・大きさを、受け取る側の領域にある命令の列で先に確保した場所に書き込みます。最後に送る側が命令を確定させると、受け取る側から見えるようになります。送る側は、データそのものはコピーしません。 受け取る側は progress の中でこの命令を読み、ハンドルを開いて、送る側の GPU のバッファを自分のプロセスから使えるようにします。開くのは cudaIpcOpenMemHandle で、開いた結果はキャッシュします。 受け取る側のスレッドが、開いたアドレスにオフセットを足した場所から自分のバッファへ、cudaMemcpy でコピーします。受け取る側は送る側の GPU のバッファから直接読み、SAR のように中継用の領域にいったん置くことはしません。 コピーの呼び出しが戻ってから、受け取る側は送る側の領域にある応答の列に処理が済んだことを書き、送る側は自分の progress でそれを見ます。送る側の書き込みが完了するのは、このときです。 1 か 2 が失敗すると、smr_do_ipc はこの書き込みを次の節の SAR に切り替えます。問題は、1 でスレッドのコンテキストの中にバッファが見つからないときに起きます。 SAR SAR では、データがホストの共有メモリを中継して 2 回コピーされます CUDA IPC を選んだあとで準備に失敗すると、shm は SAR という経路に切り替えます。SAR は segmentation and reassembly の略で、データを小さく分けて送り、受け取る側で組み立て直す方式です。切り替えるときの警告「unable to use IPC for msg, fallback to using SAR」は FI_WARN_ONCE で出るので、1 つのプロセスで 1 回しか出ません。 SAR では、受け取る側の共有メモリにある、1 つ 32 KiB の中継用の領域を使います。この領域をバウンスバッファと呼びます。送る側はメッセージごとに、空いているバウンスバッファを最大 64 個まで確保し、相手ごとの上限がある場合はその数までにします。 そのあと、2 つの手順をメッセージの終わりまで繰り返します。まず送る側が、自分の GPU のバッファから中継用の領域へデータをコピーし、データが入ったことを示すフラグを FULL にします。次に受け取る側が、そこから自分のバッファへコピーし、領域をまた使えることを示すフラグを EMPTY に戻します。GPU のメモリのデータは、ホストの共有メモリを中継する 2 段のコピーで、1 回に最大 2 MiB (64 × 32 KiB) ずつ運ばれます。SAR に切り替わることで KV 転送時間がどれだけ長くなるかは、次の節で、切り替えをなくしたときの値と比べます。 問題: 転送スレッドが向く GPU ! ここで、もとの目的に戻ります。分離推論では、prefill が作った KV cache を、decode へ渡します。この受け渡しにかかる時間は、最初のトークンが出るまでの時間 (TTFT) に入るからです。問題は、この受け渡しの送る側で起きます。prefill の TP rank 1 では、転送スレッドが別の GPU を向いていると考えられ、そのため CUDA IPC の準備に失敗し、ホストのメモリを中継する SAR に切り替わります。測った構成 (NIXL 1.1.0、shm を使う設定、page size 1、1 秒に 1 件) では、SGLang が記録する KV 転送時間の平均が、そのままで 1,886 ms、転送スレッドに正しい GPU を指定する 1 行を SGLang に足すと 461 ms でした。この 2 つは同じ回に測った組で、最後の節の表の値とは別の回の測定です。 問題: prefill の TP rank 1 では、転送スレッドが 0 番の GPU のコンテキストのまま、1 番の GPU のデータを送ろうとします CUDA の runtime API は、いま使う GPU をホストのスレッドごとに覚えていて、これを current device と呼びます。cudaSetDevice を呼ばないスレッドでは、0 番の GPU が既定になります。一方、cuMemGetAddressRange のような driver API は、そのスレッドで有効な CUDA のコンテキスト (context) の中でアドレスを探します。コンテキストは GPU ごとに作られ、その GPU で確保したメモリを管理します。 SGLang の各プロセスのメインのスレッドは、初期化のときに自分の GPU を選びます。一方、NIXL の転送を発行するスレッドは、あとから threading.Thread で作られます。このスレッドで動く transfer_worker は、GPU を選びません。 TP rank 0 の KV cache は 0 番の GPU にあるので、たまたま向きが合います。TP rank 1 では、転送スレッドが 0 番の GPU を向いたまま、1 番の GPU にある KV cache の書き込みを発行します。0 番のコンテキストの中では 1 番の GPU のアドレスが見つからず、ハンドルを作れないと考えられます。ただし、どの呼び出しで 0 番のコンテキストが有効になったかまでは確かめていません。 NIXL 1.1.0 で shm を使ったまま、起動したばかりのエンジンに 8,192 トークンの要求を 1 つ送ると、prefill の TP rank 1 のプロセスから、Failed to perform cuMemGetAddressRange: CUDA_ERROR_NOT_FOUND の警告と、SAR への切り替えの警告が出ます。切り替えた回数は、警告の行からは分かりません。回数は、最後の節で数えます。 SGLang の transfer_worker の最初に、torch.cuda.set_device(self.kv_args.gpu_id) を呼ぶ 1 行を足すと、同じ要求を送っても、どちらの警告も出ませんでした。 2026 年 10 月 19 日に確かめた時点で、この転送スレッドについての SGLang の Issue はありません。近い修正でマージ済みの sgl-project/sglang#38774 は、NIXL のバックエンドを作るスレッドで GPU を選ぶようにしたもので、main ブランチの transfer_worker は、まだ GPU を選んでいません。 直し方 1: shm を使わない設定 直し方 1 は、SGLang のプロセスを起動する前に、環境変数 FI_EFA_ENABLE_SHM_TRANSFER=0 を設定するものです。efa プロバイダは、同じノードの相手にも shm に任せず、EFA の装置で RDMA の書き込みをします。 この設定が効くのは、efa プロバイダの中に shm の経路そのものが作られないからです。問題の失敗は、shm が CUDA IPC の準備をするときに起きるので、shm を通らなければ、その準備も SAR への切り替えも起きません。SGLang の転送スレッドが GPU を選ばないことは、この設定でも変わりません。 shm を使わない設定で本当に EFA の装置を通るのかを、装置のカウンタで確かめました。rdma_write_bytes は、libfabric や NIXL の機能ではなく、EFA の装置が数えているカウンタです。装置が RDMA の書き込みで送ったバイト数を、Linux の EFA ドライバが統計の項目として登録し、カーネルの RDMA の共通部分 (RDMA core) が、それを sysfs のファイルとして見せます。ファイルの場所は、下の図のいちばん上に示しました。パスに infiniband とありますが、EFA は InfiniBand の装置ではありません。Linux の RDMA の仕組みはもともと InfiniBand のために作られたので、分類の名前が今も infiniband のままで、EFA のドライバも RDMA の装置としてここに登録されます。このファイルは装置が書き込むものではありません。読むたびに、カーネルが EFA ドライバを通して装置にその時点の値を問い合わせ、その値を返します。ただし、前の問い合わせから約 10 ms 以内に読むと、前の値をそのまま返します。 数えているのは EFA の装置で、EFA ドライバが登録し RDMA core が sysfs に出した値を、同じノードの別の Pod から読みました この記事の構成では、推論のエンジンを動かす Pod は、ホストのネットワークを使わない設定でした。この Pod の中からは、hw_counters のディレクトリが見えませんでした。同じノードに置いた、ホストのネットワークを使う Pod からは見えたので、その確認用の Pod から、負荷をかける前と後にすべてのカウンタを読んで差を取りました。負荷は、NIXL 1.1.0、page size 1 の構成に、平均して 1 秒に 1 件の割合で 15 秒のあいだ要求を送るものです。要求が届く間隔はばらつくため、この回は 15 件ではなく 17 件が届き、17 件すべてが完了しました。shm を使う既定の設定では、バイト数のカウンタはどれも 0 のままでした。shm を使わない設定では、rdma_write_bytes が 8,663,389,312 バイト増えました。 送るはずの KV cache の量は、モデルの設定から計算すると 8,663,334,912 バイトで、測った値との比は 1.0000063 です。負荷をかけない 20 秒にはカウンタが動かなかったことと合わせると、この測定の区間の KV cache の転送は、装置の RDMA の書き込みとして出たと読めます。差の 54,400 バイト (17 件で割ると 3,200 バイト) の内訳は分かっていません。 shm を使わない設定で起動したエンジンのログには、EFA will not use SHM for intranode communication because FI_EFA_ENABLE_SHM_TRANSFER=0 が出て、shm の行は 1 行も出ませんでした。shm の経路が作られないことそのものは、上で示したコードの分岐から言えることで、shm の中を直接見て確かめたわけではありません。 測り方と計算 項目 中身 数えたもの EFA の装置の hw_counters の全項目。rdma_write_bytes、rdma_write_recv_bytes、rdma_write_wrs、tx_bytes など 読んだ場所 同じノードに置いた、ホストのネットワークを使う確認用の Pod から sysfs を読んだ 負荷 入力 8,192 トークン・出力 256 トークンの要求を、平均 1 秒に 1 件の割合で 15 秒。届く間隔は指数分布でばらつき、この回は 17 件が届いて 17 件とも完了。NIXL 1.1.0、page size 1、TP 2 雑音の確認 どちらの設定でも、測る直前に負荷をかけない 20 秒のあいだ、カウンタはどの項目も動かなかった 送るはずの量 このモデルは MLA という Attention の方式で、KV を 1 層、1 トークンあたり 576 個の値に圧縮して持つ。576 は、圧縮した KV の 512 と位置の情報の 64 の和。1 トークンあたり 27 層 × 576 × 2 バイト = 31,104 バイト。1 件、1 rank で × 8,192 トークン = 254,803,968 バイト。MLA の KV は TP の 2 つの rank が同じものを持って送るので、× 2 rank × 17 件 = 8,663,334,912 バイト 使った装置 prefill と decode のどの rank も、Pod に割り当てた 1 台の装置を開き、その装置の中で折り返した。ノードのほかの装置のカウンタは 0。割り当てられる装置は Pod ごとに変わりうるので、ノードのすべての装置のカウンタを読んだ shm を使わない設定では、同じノードの中のほかの相手への libfabric の通信も、shm ではなく EFA の装置を通ります。 直し方 2: NIXL 1.4.0 直し方 2: NIXL 1.4.0 は、転送を発行するスレッドのコンテキストを、自分で合わせ直します NIXL 1.4.0 で問題が起きないのは、次の順に、TP rank 1 の転送スレッドが 1 番の GPU のコンテキストを向くようになるからです。 メモリを登録するとき、NIXL は KV cache のバッファを持っているコンテキストを調べて、エンジンに 1 つ記録します。TP rank 1 では、1 番の GPU のコンテキストです。 GPU のメモリから書き込みを出す前に、NIXL は、書き込みを出すスレッドを、記録したコンテキストに切り替えます (cuCtxSetCurrent)。これは既定の設定で、1 つのプロセスが 1 つの GPU だけを登録している場合の動きです。書き込みを出すのは、NIXL を呼んだ SGLang の転送スレッドです。 そのあと同じスレッドの上で、shm が IPC ハンドルを作る準備の中で、libfabric の cuMemGetAddressRange が呼ばれます。スレッドのコンテキストが KV cache のある GPU と合っているので、測った構成では、この呼び出しが失敗しなくなりました。 ハンドルの準備が失敗しないので、SAR への切り替えの分岐に入りません。 NIXL 1.1.0 も、登録のときに同じようにコンテキストを記録します。ただし、それを設定するのは登録の処理の中だけで、書き込みを出す前には設定し直しません。書き込みの前に設定し直す処理は、ai-dynamo/nixl#1581 の変更に含まれています。直ったのは NIXL の側で、SGLang の転送スレッドのコードは変わっておらず、自分では GPU を指定しないままです。 この動きは既定で有効です。ただし、登録のときに 1 つのプロセスが複数の GPU を使うと分かった場合は、代わりに書き込みのたびに cudaSetDevice で GPU を合わせます。 測った構成では、各プロセスは自分の GPU だけを登録するので、複数の GPU を使うという判定にはならず、ログの「Multi-GPU detected」は 0 行でした。働いたのは、登録時のコンテキストを設定し直す動きのほうです。これにより、SGLang を直さなくても、TP rank 1 でハンドルを作れるようになりました。実際に経路がどうなったかは、最後の節で数えます。 直し方 2 では、KV cache は shm の経路のまま、CUDA IPC で運ばれます。 4 つの組み合わせの値 最後に、NIXL の版と shm の設定の 4 つの組み合わせで測った値を並べます。元の構成は NIXL 1.1.0 で shm を使う設定、直し方 1 は NIXL 1.1.0 で shm を使わない設定、直し方 2 は NIXL 1.4.0 で shm を使う設定、両方は NIXL 1.4.0 で shm を使わない設定です。どれも page size 1 です。 組み合わせ KV 転送時間 ms 経路 TTFT の中央値 ms 元の構成 1,776 rank 0 が CUDA IPC、rank 1 が SAR 2,453 直し方 1 196 EFA 1,376 直し方 2 533 CUDA IPC 2,410 両方 267 EFA 1,389 負荷は、どれも平均して 1 秒に 1 件です。KV 転送時間は、SGLang が要求ごとに記録する prefill_transfer_kv_cache を、測ったあいだの全件で平均した値です。1 件ごとに、prefill 側が KV cache の送信を頼んでから、送り終わりを確認するまでの時間で、確認はスケジューラの1 周に 1 回なので、送り終わってから気づくまでの待ちも含みます。TTFT は、利用者から見た、要求を送ってから最初のトークンが返るまでの時間で、KV 転送時間はその一部です。表の TTFT は、20 秒の負荷を 3 つの seed で流して合わせた 66 件の中央値です。行ごとに別の回の測定です。直し方 2 では、KV 転送時間は元の構成より短い一方、TTFT の中央値はほぼ同じでした。その理由は、この記事では調べていません。 経路の列は、prefill から decode へ KV cache を運んだ経路です。表の時間を測った回とは別の回に、4 つの組み合わせのそれぞれで、NIXL の版と shm の設定だけを変え、page size 1、平均 1 秒に 1 件の負荷を 15 秒かけて判定しました。どの回も 17 件の要求が届き、計算した転送の量は「直し方 1」の節と同じです。判定には 3 つのものを使いました。1 つ目は EFA の装置のカウンタ rdma_write_bytes の増分で、計算した量との比を見ました。2 つ目は、送る側で IPC ハンドルを作る cudaIpcGetMemHandle と、受け取る側の CUDA IPC のコピーの回数です。3 つ目は、送る側の cuMemGetAddressRange の失敗と、SAR の中継用の領域を通るコピーの回数です。2 つ目と 3 つ目は、bpftrace で、プロセスごとに呼び出しを数えました。 元の構成では、装置の増分は 0 でした。prefill の rank 0 では cudaIpcGetMemHandle が 5,454 回呼ばれ、decode の rank 0 が CUDA IPC で 5,454 回、合わせて 4,331,667,456 バイトをコピーしました。これは rank 1 つ分の KV cache の量 (254,803,968 バイト × 17 件) と同じです。prefill の rank 1 では cuMemGetAddressRange が 5,454 回失敗し、中継用の領域を通るコピーが 136,782 回で、運んだ量は rank 0 と同じでした。SAR への切り替えの警告は、SAR の節で見たとおりプロセスごとに 1 回しか出ないので 1 行だけでしたが、切り替えが起きた回数は失敗の回数と同じ 5,454 回です。直し方 2 では、装置の増分は 0 で、prefill の rank 0 と rank 1 のどちらでも cudaIpcGetMemHandle が 1,998 回呼ばれ、decode の rank 0 と rank 1 がそれぞれ 1,998 回、rank 1 つ分の KV cache と同じ量を CUDA IPC でコピーしました。cuMemGetAddressRange の失敗と SAR への切り替えの警告は 0 行でした。直し方 1 の装置の増分は、「直し方 1」の節で見たとおりです。両方でも、装置の増分は直し方 1 と同じバイト数で、4 つのプロセスのどれでも IPC ハンドルの作成とコピーは 0 回で、SAR への切り替えの警告も 0 行、装置の RDMA の書き込みは 4,066 回でした。 まとめ この記事では、同じノードの相手へ KV cache を送るときに、libfabric の efa プロバイダと shm プロバイダがどう振る舞うかを紐解きました。efa プロバイダは既定で同じノードの相手を shm プロバイダに任せ、shm プロバイダは条件がそろうと CUDA IPC を選びます。送る側のスレッドのコンテキストにバッファが見つからないと SAR に切り替わり、FI_EFA_ENABLE_SHM_TRANSFER=0 では EFA の装置を通ります。流すデータの量や条件で結果は変わるので、この記事の数値は目安でしかありませんが、この記事の構成で、どの経路をいつ通るかという挙動は、4 つの組み合わせで確かめられました。