0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

WSL2 では cudaHostRegister したメモリに GPU カーネルが書けない — sglang HiCache が hybrid GDN で落ちる真因

0
Last updated at Posted at 2026-09-21

2枚の12GB GPUで27Bモデルを128K文脈で回している環境に、sglang の HiCache(KVキャッシュをホストRAMへ退避する機能)を入れようとしたら、最初の長文リクエストで必ず落ちました。

CUDA error: an illegal memory access was encountered

真因は2つあり、どちらも量子化KV(int2)とは関係なく、片方は WSL2 固有でした。同じ構成で詰まっている人向けに、切り分けの手順ごと残します。

結論を先に

  1. WSL2 では cudaHostRegister したホストメモリに GPU カーネルが書き込めない。 フラグ(Default/Portable/Mapped/両方)を全部試しても駄目で、cudaHostAlloc(PyTorch の pin_memory=True)なら通る
  2. sglang の HiCache ホストプールは既定で cudaHostRegister を使う。mamba(GDN状態)のバックアップカーネルはホストポインタへ GPU から直接ストアするので、WSL2 では必ず落ちる
  3. 量子化KVを使っている場合は別の問題も重なる。ホスト側の1トークン幅が3か所で食い違っていた

環境

GPU 12GB × 2(TP2)
ホスト 32GB RAM、WSL2
推論エンジン sglang 0.5.19
モデル 27B dense、hybrid GDN + full attention(64層のうち16層が full attention、48層が mamba)
KVキャッシュ int2 量子化(packed 2bit + グループごとの scale/zero)

症状と、最初に間違えたこと

--enable-hierarchical-cache を付けると、最初の長文リクエストでスケジューラごと死にます。Python のトレースバックは出ますが、指している行は自分が書いた scales/zeros の転送でした。

File ".../pool_host/mha.py", line 1274, in _transfer_int2_scales_zeros
    h[h_idx] = d[d_idx].to("cpu")
torch.AcceleratorError: CUDA error: an illegal memory access was encountered

ここで「自分のコードが悪い」と決めつけると時間を溶かします。実際、同じメッセージが親切に断っています。

CUDA kernel errors might be asynchronously reported at some other API call,
so the stacktrace below might be incorrect.

CUDA のエラーは非同期です。この行は、非同期エラーを回収した最初の同期点でしかありませんでした。

切り分け1: CUDA_LAUNCH_BLOCKING=1

これを付けて再起動すると、発生源が変わります。

File ".../pool_host/mamba.py", line 465, in backup_from_device_all_layer
  File ".../kernels/ops/mamba/transfer_mamba.py", line 75, in transfer_kv_mamba_lf_pf
RuntimeError: Failed at .../kvcacheio/transfer_mamba.cuh:184:
              CUDA error: an illegal memory access was encountered

int2 の KV 転送ではなく、mamba の状態をホストへ退避するカーネルでした。

切り分け2: カーネルに渡っているものを全部照合する

推測で直すと外します(3回外しました)。カーネルが見ているのは「層ごとの先頭ポインタの配列」と「2つのインデックス列」だけなので、発射直前に全部突き合わせるコードを入れました。

  • キャッシュしたポインタ vs 現在のテンソルの data_ptr()(stale pointer の検出)
  • インデックスの範囲(デバイス側スロット数・ホスト側スロット数)
  • ストライド(stride(0) * element_size == item_size か、is_contiguous() か)
  • ホスト側1スロットの幅 == item_size * num_layers か

全部通りました。 そして実測値がログに出ます。

num_layers=48 item_size=786432B
device[shape=(33, 24, 128, 128) slots=33]
host[shape=(65, 48, 1, 24, 128, 128) slots=65 stride0=37748736B]
src_idx=[6,6] dst_idx=[0,0] n=1

幅もスロット数も整合しています。つまり入力は正しい。

切り分け3: カーネル単体に落とす

サーバの起動に毎回4分かかるので、同じ形状でカーネルだけを叩く再現スクリプトを書きました。30秒で再現します。ここで初めて核心に触れました。

dst(書き込み先のホストバッファ)の確保方法だけを変えて比較します。

dst の確保方法 結果
GPU 上 OK
torch.zeros(..., pin_memory=True) = cudaHostAlloc OK
アロケータ + cudaHostRegister(ptr, size, flags) illegal memory access
素の CPU テンソル(未ピン) FAIL(当然)

さらにフラグを振ります。

flags 結果
0 (cudaHostRegisterDefault) FAIL
1 (Portable) FAIL
2 (Mapped) FAIL
3 (Portable|Mapped) FAIL

全滅。 cudaHostAlloc だけが通ります。

CUDA のドキュメントでは、UVA のあるシステムなら cudaHostRegisterDefault でも mapped かつ portable になる、と書かれています。WSL2 ではその前提が成り立ちません。 malloc したホストメモリは、登録してもデバイスのアドレス空間に入らないようです。

なぜ KV 側(MHA)は落ちなかったかというと、page_first レイアウトではデバイス上のステージングバッファを経由する経路を通るためでした。mamba だけがホストへ直接書いていた、という非対称です。

直し方

ホストプールの確保関数を差し替えます。WSL を検出したら cudaHostAlloc 経路にする、環境変数で上書きできる、という形です。

def _host_register_is_mappable() -> bool:
    override = os.environ.get("SGLANG_HICACHE_HOST_ALLOC", "").strip().lower()
    if override == "pin":
        return False
    if override == "register":
        return True
    return "microsoft" not in platform.uname().release.lower()

注意点として、cudaHostAlloc 経路は mmap/shm 系のアロケータと併用できません。既定以外を要求されていたら黙って無視せず落とすようにしています。

おまけ: エラーメッセージ自体が壊れていた

切り分けの途中で、登録失敗時のメッセージ生成がそもそも例外を投げているのを見つけました。

raise RuntimeError(
    f"cudaHostRegister failed (rc={rc}, "
    f"{cudart.cudaGetErrorString(rc)}) ..."   # rc は int
)

cudaGetErrorString() は torch._C._cudart.cudaError しか受け取りません。int を渡すと TypeError になり、本来の失敗理由が握り潰されます。こういうのは、踏まないと見つかりません。

もう1つの真因: 量子化KVのホストプールの幅

int2 を使っている場合はこちらも必要です。素の MHATokenToKVPoolHost は1トークンを head_num * head_dim * itemsize バイトの連続領域として扱いますが、int2 では2点で成り立ちません。

  1. packed codes は [slots, head_num, head_dim // 4] の uint8。1ヘッドあたり head_dim ではなく head_dim // 4 バイトしかない
  2. scale/zero は k_scales_zeros という別テンソル。運ばないと、復元したスロットは前の占有者の量子化パラメータで解釈され、静かに壊れる

最初の実装では「1トークン = packed + sz の連続領域」とみなして head_dim を実効バイト数に差し替えたのですが、これが誤りでした。sz は結局別バッファに置いたので、本体だけが幅広になっていたのです。

本番実値(head_dim=256 / head_num=2 / group_size=64 / scale_dtype=float32)で並べるとこうなります。

旧 新
デバイスの1行 128B 128B
ホストの token_stride_size 192B 128B
element_dim * itemsize 512B 128B

3つが三者三様でした。element_dim は親クラスがデバイスの head_dim から算出するため、そこだけ4倍幅のまま残っていたのです。

直し方は「ホスト本体は packed だけを持ち、sz は別バッファで同じインデックスで運ぶ」。そのうえで起動時に3つの幅を突き合わせる assert を入れました。食い違ったら起動させません。この手の食い違いは、実行時には必ず illegal memory access として出てきて、原因の特定に非常に時間がかかるからです。

Int2 HiCache host pool geometry verified: 256 B/token/layer packed (K+V),
128 B/token/layer scales+zeros (K+V), 16 layers, 60001 slots.

動いた。次はホストRAMの壁

起動して長文も通るようになったので --hicache-ratio 2(既定値)にしたら、今度は起動しません。

ValueError: Not enough host memory available.
Requesting 2.55 GB but only have 0.52 GB free.

sglang はホストプールの上限を (MemAvailable - 10GiB) / ranks_per_host で決めます。TP2 なので実質半分です。ここで2つ、実機でないと分からないことがありました。

ホストプールは「モデル重みのロード中」に確保される。 つまり MemAvailable が最も低い瞬間です。32GB のマシンで実測すると、ロード中に 565MB まで落ちていました。free -m を平常時に見て判断すると外します。

上流に既知のバグがある。 issue #38156 / PR #38157(未マージ)で、各ランクが別々のタイミングで available を読むのに常にランク数で割るため、先行ランクのプールが二重に差し引かれます。TP群で MIN の all_reduce を取る修正を移植しました。

ただし今回のケースでは、二重計上ではなく本当に足りていませんでした。ログの順序が証拠です。

hierarchical KV host pool: 420001 tokens, 2.58 GB   ← KV は両ランクとも成功
Not enough host memory available. Requesting 2.55 GB but only have 0.59 GB free
                                                    ← 次の mamba で落ちた

ratio を上げると何が起きるか

予約を下げれば ratio 2 でも起動します。物理的には入ります。しかし。

--hicache-ratio ピン留め(2ランク計) MemAvailable 実効ページキャッシュ
1 5.2 GB 10.2 GB 2.0 GB
2 10.3 GB 3.7 GB 589 MB

ratio 2 にすると、アイドルなのにスワップインが継続的に発生し、PC 全体が目に見えて重くなりました。

si=547  si=172  si=0  si=376  si=68  (KB/s)

理由は単純で、ピン留めしたメモリはスワップに逃がせないからです。GPU が DMA で触るため物理ページを固定する必要があり、足りなくなると固定できない側(他のプロセス)だけが押し出されます。ページキャッシュも潰れるので、ファイルアクセスが毎回ディスクに行きます。

32GB のホストでは ratio 1 が実用上の上限でした。なお mamba のホストプールは1スロット 37.7MB と非常に重く、33スロットで 1.25GB です。210,001トークンぶんの int2 KV とほぼ同じコストになります。

効果(実測)

同じビルドで HiCache の ON/OFF を比較しました。

OFF ON (ratio 1)
冷prefill(75Kトークン) 1010 tok/s 964 tok/s
decode(bs1) 51.5 tok/s 50.7 tok/s
短文TTFT 0.132–0.153 s 0.123–0.153 s
長文prefill中の短文TTFT 0.140 s 0.141 s
並列8・別プレフィクス・2周目 5.2 s 1.6 s

再訪が 3.2倍速。 別プレフィクス8本という最も重い条件で、デバイスプールから溢れた分をホストから戻せています。代償は冷prefill −4.5%、decode −1.6% で、write_through により毎 insert でホストへ書くコストです。

品質も確認しました。96K文脈での集計・順序列挙タスク20問で、ON/OFF が設問単位まで完全一致(18/20)。壊れていないことは確認できています。

使えないもの

--hicache-mem-layout layer_first は MambaPoolHost が assert で拒否します。page_first を使ってください。これは起動即死なので分かりやすい方です。

L3(--hicache-storage-backend によるディスク退避)は int2 では使えません。 書き出される flat data page が packed codes だけで scales/zeros を含まないため、復元時に前の占有者の量子化パラメータで解釈されて静かに壊れます。NotImplementedError で塞いであります。「ホストRAMが足りないならディスクへ」が効かないのは、ここが理由です。

踏んだ落とし穴のまとめ

  • CUDA の illegal memory access は、トレースバックが指す行が発生源とは限らない。 まず CUDA_LAUNCH_BLOCKING=1
  • 推測で直さない。 カーネルに渡っている値を全部 assert で突き合わせてから考える
  • サーバ起動が重いなら、カーネル単体の再現に落とす。 4分 → 30秒になり、条件を振れるようになる
  • ログの oom 文字列検索は bootstrap_room=None に誤ヒットする。 out of memory / oom-kill で見る
  • ホストプールの確保タイミングは「重みロード中」。free -m の平常値では判断できない
  • ピン留めメモリは他人を追い出す。 RAM に「入るか」ではなく「固定していいか」で考える

成果物

修正は sglang のフォークに、v0.5.19 タグ起点のブランチとして置いてあります。base ブランチが素の v0.5.19 なので、PR の差分がそのまま追加分になります。

GPU 不要で回せる検証スクリプトを2つ同梱しました。

  • int2_hicache_geometry_test.py — ホストプールの3つの幅が一致するか。旧バグを復元すると落ちる負のテスト付き
  • mamba_transfer_repro.py — このマシンで cudaHostRegister したメモリがデバイスからマップされるか

WSL2 で HiCache が謎に落ちる人は、まず後者を回してみてください。30秒で白黒つきます。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?