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?

配布wheelにsglangのPRを自力で統合する — Escha-W2(2bit MoE)にOSCAR INT2 KVを載せて262Kコンテキストを出すまで

0
Last updated at Posted at 2026-09-15

この記事の統合手順は、dense版(Qwen3.8-27B)にも対応した「手順書」版にまとめ直しました。 コピペで進められる形になっていて、ファイル単位の対応表・回転行列の生成スクリプト・統合チェックスクリプト・dense/MoEの差分一覧まで入っています。これから作業する方はそちらを推奨します。

👉 【手順書】dense でも MoE でも通る — sglang の未マージPR(OSCAR INT2 KV)を Escha-W2 ランタイムに統合する

(この記事は、なぜ git apply が通らないのか・較正した回転行列が逆効果だった話など、経緯と判断の記録として残しています)

この記事は、以前の記事(RTX 5070×2(12GB×2)でQwen3.8-27B-Escha-W2を動かす)の第2部にあった誤りの訂正を兼ねています。

あちらでは「実際に必要だったコード修正」を4点として書きましたが、実際の統合は10ファイルに及びます(コピー3本 / 手マージ6本 / 移植不要6本)。また「MoE版にはヘルパー関数3つが移植漏れしていて、dense版の decode_attention.py からコピーが必要」と書いた箇所は誤りでした。素のwheelを展開して確認したところ、_get_scale_group_size / _get_shared_kv_scale_group_size はPR #32129側のヘルパーで、_safe_block_h は私が書いた関数です。dense版・MoE版どちらのwheelにも、OSCAR関連のコードは一切入っていません。

この記事は、その統合作業をファイル単位の対応表まで含めて書き直したものです。dense版(Qwen3.8-27B)との差分は第6部にまとめました。

12GBのコンシューマGPU 2枚で 35B MoE を 262,144トークンのコンテキスト・8人同時利用 で回しています。成立させている要素は2つで、どちらも配布物としては別々に存在しています。

  • Escha-W2: Qwen3.6-35B-A3B を2bit量子化したMoE。ランタイム(sglangのフォーク)がwheelで配布される
  • OSCAR INT2 KV: KVキャッシュを2bit化する実装。sglang本家の未マージPR #32129にしかない

この2つは一緒に配られていないので、wheelの中のsglangにPRを手で統合する必要があります。この記事はその統合手順そのものの話です。設定値のチューニングではなく、「他人のフォークに他人のPRを載せる」作業の具体的な手順書です。

結論から言うと、git apply は通りません。通らない理由と、代わりに何をどこへ持っていけばいいかの対応表を全部載せます。

前提と所要時間

正直に書きます。ステップ1・2・4・5は普通のインストール作業で30分程度、ステップ3(統合)だけで半日かかります。スクリプトで自動化できる作業ではありません。CUDAカーネルを書ける必要はありませんが、diffを読んで関数をファイル間で移す作業に抵抗が無いことは必要です。

GPU 12GB × 2(compute capability 12.0 / RTX 50系)
ディスク モデル12GB + sglangのチェックアウト
Python CPython 3.12 固定(wheelが cp312 のみ)
glibc 2.28以上(manylinux_2_28)
CUDA toolkit 不要(ptxas は triton に同梱)
OS WSL2 / Ubuntu(素のLinuxでも可。LD_LIBRARY_PATH の行だけWSL固有)

合計24GB未満のVRAMではこの構成は入りません。多いぶんには問題なく、後半のチューニングもシビアでなくなります。

また、wheelは sglang という名前のパッケージを入れますが本家sglangではありません。同じ環境に本家sglangを入れると壊れるので、この環境はランタイム専用にしてください。

なぜKVキャッシュを2bitにするのか

262,144トークンぶんのKVキャッシュに必要なVRAMです(GPU 1枚あたり、実測):

KV dtype 必要VRAM
fp16 2.52 GB
INT2(OSCAR) 0.39 GB

モデル実体が12GB、GPUが12GB×2という構成なので、この2.1GBの差がそのまま「262Kが載るか載らないか」になります。KV量子化は長コンテキストの高速化テクニックではなく、そもそも入るかどうかの前提条件です。

最終的な実測値:

コンテキスト長 262,144
KVプール 350,257トークン(全リクエスト共有)
KVサイズ K 0.21 GB + V 0.21 GB
VRAM 11.3 / 11.7 GB(12GB中)
単独デコード 87.5 → 56.8 tok/s
8並列・400トークン応答 集約 295 tok/s
180Kコールドprefill 46.2秒
同じ180Kを再送(prefix cacheヒット) 1.87秒

第1部: なぜ git apply が通らないのか

最初は「PRブランチのソースに差分を当てたフォークを作る」つもりで始めて、途中で断念しました。理由は2つです。

理由1: ツリーのレイアウトが違う

wheelの中のsglangと、PRが書かれている本家の現在のツリーは、ファイル配置が別物になっています。決定的だったのがINT2デコードカーネルです。

パス
PR側 python/sglang/kernels/ops/attention/decode_attention.py
wheel側 sglang/srt/layers/attention/triton_ops/decode_attention.py

パッケージが違うので、ファイルを移動させるだけでは import が解決しません。ほかにも scheduler.py は1ファイルで3,800行以上の差があります。共通の祖先が無い2つのツリーだと思ってください。

理由2: 触るファイルの半分以上が「既にescha側が書き換えたファイル」

PRが触るファイルのうち、wheelに存在するものを実際に数えると10本。その内訳が重要です。

  • 4本: wheelに存在しない = そのままコピーできる
  • 6本: escha側が既に改造済み = 上書きするとランタイムが壊れる。手でマージするしかない

「PRのファイルを入れてください」の一文で済ませられないのはここです。以下、その対応表です。


第2部: 統合の対応表

ステップ0: PRをコミット固定で取得する

PRは動きます。後から手順が腐らないよう、取得したコミットを記録してください。

git clone https://github.com/sgl-project/sglang.git sglang-pr32129
cd sglang-pr32129
git fetch origin pull/32129/head:oscar-int2-kv
git checkout oscar-int2-kv
git rev-parse HEAD   # 本記事の対応表は 4b28b26 時点のもの

インストール先(site-packages)はこう取ります。

SP=$(python -c "import sglang, pathlib; print(pathlib.Path(sglang.__file__).parent.parent)")

ステップ1: そのままコピーできる3本

いずれもwheel側に存在しないファイルなので、パスを変えずにコピーするだけです。

for f in sglang/QuantKernel/oscar_rotation_clip_int2_kv.py \
         sglang/srt/layers/attention/quantized_kv_prefill.py \
         sglang/srt/mem_cache/kv_quant_kernels.py; do
    cp "python/$f" "$SP/$f"
done

ステップ2: 手でマージする6本

ここが本番です。escha側のファイルをベースに、PR側のOSCAR関連の追加分だけを移植します。逆(PR側をベースにする)は絶対にやらないでください。escha固有の改造が消えます。

ファイル 差分規模 移植するもの
srt/environ.py 12行 環境変数の定義: SGLANG_OSCAR_{K,V}_ROTATION_PATH、SGLANG_OSCAR_{K,V}_CLIP_RATIO、SGLANG_LLOYD_MAX
srt/server_args.py 11行 --kv-cache-dtype の choices に int2 を追加。新規引数 --kv-cache-quant-group-size とフィールド
srt/model_executor/model_runner.py 5行 kv_cache_dtype == "int2" の分岐
srt/model_executor/model_runner_kv_cache_mixin.py 52行 KVプール生成のint2分岐、kv_cache_quant_group_size のプールへの受け渡し
srt/mem_cache/memory_pool.py 411行 OscarRotationConfig、load_oscar_rotation_config()、load_oscar_rotations()、_resolve_quant_grouping()、scale/zeroバッファ一式(_allocate_scales_zeros_buffers、get_{key,value}_scales_zeros、get_raw_{key,value}_buffer、get_oscar_rotation)、dtype == "int2" の確保パス
srt/layers/attention/triton_backend.py 292行 _forward_extend_int2() とそのディスパッチ(is_int2 = getattr(kv_pool, "dtype", None) == "int2")、_apply_oscar_rotation、apply_inverse_v_rotation
srt/layers/attention/triton_ops/decode_attention.py 1,494行 INT2デコードカーネル一式。PR側の kernels/ops/attention/decode_attention.py から持ってくる

差分規模は「稼働中のインストールと素のwheelのdiff行数」で、第3部のパッチぶんも含んだ数字です。

一番重いのは decode_attention

移植するシンボルはこれだけあります。

_fwd_kernel_stage1_quant_int2
_fwd_grouped_kernel_stage1_quant_int2
_decode_att_m_fwd_quant_int2
_decode_grouped_att_m_fwd_quant_int2
decode_attention_fwd_normal_quant_int2
decode_attention_fwd_grouped_quant_int2
decode_attention_fwd_quantized
_get_scale_group_size / _get_shared_kv_scale_group_size

PR側にある decode_attention_fwd_int2_unified(ディスパッチャ)は移植不要でした。wheel側は triton_backend.py から2つのエントリポイントを直接呼ぶ構造だからです。パッケージが違うので、ファイルを丸ごと持ってくるのではなく、import と周辺ヘルパ名を移植先のツリーに合わせて書き直す作業になります。

ステップ3: 移植しないもの

PR #32129 は他にも srt/mem_cache/kv_cache_dtype.py、srt/mem_cache/unified_kv_pool.py、srt/models/{qwen3,glm4_moe,utils}.py、QuantKernel/gpu_flush_int2.py を触ります。

  • 前2つはwheel側のツリーに存在しません(sglangの新しいKVプール抽象。eschaのフォークはそこまで追随していない)
  • モデル定義はEscha-W2には不要
  • プール内部にあるデータ不要Hadamardフォールバックも移植していません

最後の点が次の節に効きます。フォールバックが無い = 回転行列のパス指定が必須になります。


第3部: 回転行列を作る

OSCARはKVを2bit化する前に head_dim 方向の回転をかけます。この回転行列を .pt で渡す必要があるのですが、較正(calibration)は不要です。中身はただのHadamard行列でした。

実際に本番で使っているファイルの構造:

{"layers": {0: {"rotation": tensor([256, 256], dtype=float32)},
            1: {"rotation": ...},
            ...  # 層ごとに1枚。較正なしなら中身は全層同一
            }}

必要な数字は head_dim と層数の2つだけで、どちらもモデルの config.json にあります。このチェックポイントでは text_config の下に入れ子になっているので注意してください。

python -c "
import json; c=json.load(open('/path/to/model/config.json'))
c=c.get('text_config',c); print(c['head_dim'], c['num_hidden_layers'])"
# -> 256 40

ここで2点、直感に反するが正しいことがあります。

  • 層数ぶん作るが、実際に使われるのは一部の層だけ。これはハイブリッドモデルで、config.json の layer_types を見ると4層に1層だけが full_attention(40層中10層)です。KVキャッシュを持つのはその層だけ。参照は層インデックスなので、全層ぶん作るのが正解で、多すぎるぶんは無害です(本番で使っているファイルは40層のモデルに対して64層ぶん入っていますが、問題なく動いています)
  • 全層が同じ行列。較正なしの回転は構成上どの層でも同じになります。較正した回転は層ごとに違うので、.pt を開けばどちらなのか一目で分かります

中身は Sylvester構成のHadamard行列を sqrt(256) = 16 で割ったもので、直交誤差は厳密に 0.0 でした。GPUも重みも較正データも要りません。

import math, pathlib, torch

def hadamard(n: int) -> torch.Tensor:
    h = torch.ones(1, 1)
    while h.shape[0] < n:
        h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0)
    return h

head_dim, num_layers = 256, 40   # config.json の text_config から
rot = (hadamard(head_dim) / math.sqrt(head_dim)).to(torch.float32)
assert (rot @ rot.T - torch.eye(head_dim)).abs().max().item() == 0.0

state = {"layers": {i: {"rotation": rot.clone()} for i in range(num_layers)}}
out = pathlib.Path("/path/to/oscar_rotations"); out.mkdir(parents=True, exist_ok=True)
for name in ("k", "v"):
    torch.save(state, out / f"{name}_rotation_hadamard_hd{head_dim}.pt")

このスクリプトの出力が、本番で使っているファイルとビット一致することを確認しています。

export SGLANG_OSCAR_K_ROTATION_PATH=/path/to/oscar_rotations/k_rotation_hadamard_hd256.pt
export SGLANG_OSCAR_V_ROTATION_PATH=/path/to/oscar_rotations/v_rotation_hadamard_hd256.pt

較正した回転行列は「悪化」した

OSCARの論文は較正なしの劣化を大きく報告しています(Qwen3-32Bで-66.29pt)。ところがこのチェックポイントでは逆でした。長文の識別テストで:

スコア
較正なし(素のHadamard) 40/41
較正あり 37/41

出現順の完全列挙タスクの再現率が 0.969 → 0.819 に落ちました。論文の数字を信じて較正版を採用していたら、静かに品質が落ちたまま運用していたことになります。較正したら必ず測る。

ついでに、この差は needle-in-a-haystack では一切検出できませんでした(どちらも15/15)。識別に必要だったのは「似た紛らわしい情報を多数混ぜる」「複数位置の情報を集約させる」「順序どおりに列挙させる」の3点です。


第4部: パッチと落とし穴

統合が終わった上に、本番で必要だった修正を3系統当てています(公開しています。リンクは末尾)。

パッチ 効果
INT2 prefillの一時メモリ削減(3ファイル) 192Kで −746MB(−37%)。展開したprefixを3重に持っていた(3P+2E)のを、まとめて確保して out= で直接dequantする(P+E)。長コンテキストと8並列を両立させている要因
チャンク長の抑制(2ファイル) 長文prefill中に他のストリームの生成が止まる問題。トークン間隔 5.1秒×21回 → 0.14秒。既定値0で従来と完全に同一挙動
_safe_block_h(decode_attention) 静かな正しさのバグ。下記

_safe_block_h: assertもNaNも出ないバグ

grouped decodeのカーネルはヘッドをこう見ています。

VALID_BLOCK_H = min(BLOCK_H, kv_group_num)
cur_head      = cur_head_id * VALID_BLOCK_H + arange(BLOCK_H)
cur_kv_head   = cur_head_id // cdiv(kv_group_num, BLOCK_H)

cur_head はフラットなクエリヘッド番号、cur_kv_head はブロック番号由来です。両者が一致するのはヘッドブロックが1つのKVグループに収まるとき、つまり BLOCK_H >= kv_group_num または kv_group_num % BLOCK_H == 0 のときだけ。

本家は BLOCK_H = 16 をハードコードしているので2冪の kv_group_num では安全で、だから今まで表面化していません。ただし kv_group_num > 16 かつ16の倍数でない構成(96クエリヘッド / 4 KVヘッド = 24 など)では到達可能で、BLOCK_H をバッチサイズ依存にした瞬間に踏みます。BLOCK_H=4 / kv_group_num=6 なら、ヘッドブロック1がクエリヘッド4..7を覆いつつ cur_kv_head=0 を名乗るので、クエリヘッド6と7が別のKVヘッドのキャッシュを読みます。

assertも落ちず、NaNも出ず、ベンチのスコアにしか出ません。_safe_block_h は2冪に切り上げて VALID_BLOCK_H == kv_group_num を保証します。

落とし穴1: 匿名化した .patch は -p0 では当たらない

公開用にdiffヘッダの実パスを <site-packages>/... に置換したところ、patch -p0 が1本も当たらなくなりました。

$ patch -p0 --dry-run < foo.patch
can't find file to patch at input line 3
No file to patch.  Skipping patch.

-p0 はヘッダのパスをそのまま探すので、<site-packages> というディレクトリを探しに行って失敗します。しかも patch はエラー終了せずスキップして先に進むので、ループで回していると「全部当たったつもり」になります。正解は -p1(先頭1階層を剥がす)。

cd "$SP"
for p in /path/to/patches/*.patch; do
    patch -p1 --batch --backup --suffix=.orig < "$p"
done

自分で公開する側になる場合は、匿名化した後に必ず当て直しのdry-runをしてください。私はこれを公開してから気付きました。

落とし穴2: wheelを入れ直すと全部消える

統合もパッチも site-packages を直接書き換えています。pip install し直した時点で消えます。.orig バックアップを残すか、修正版を PYTHONPATH のオーバーレイに置いてください。

落とし穴3: 動いた ≠ INT2が効いている

チャットが返ってきても、INT2パスを通っている保証にはなりません。確認は2点。

起動ログ:

KV Cache is allocated. #tokens: 350257, K size: 0.21 GB, V size: 0.21 GB
max_total_num_tokens=350257, ... max_running_requests=12, context_len=262144

ここがGB単位の大きな数字なら、--kv-cache-dtype int2 が黙って無効化されています。

統合の有無:

from sglang.srt.mem_cache.memory_pool import load_oscar_rotation_config
from sglang.srt.layers.attention.triton_ops.decode_attention import (
    decode_attention_fwd_grouped_quant_int2,
)
print("OSCAR INT2 present:", load_oscar_rotation_config())

ImportError ならステップ1/2が未完。SGLANG_OSCAR_K_ROTATION_PATH を名指しする ValueError なら統合は成功していて環境変数だけが未設定です(フォールバックを移植していないので、サーバも同じ理由で起動を拒否します)。


第5部: 起動設定で外してはいけないもの

統合が終わった後の話ですが、どれも壊して学んだものです。

設定 理由
--ep-size 2 --tp-size 2 だけだとexpert重みが両GPUに複製され、12GBに対し11.06GBでロードできない
ATTN_BACKEND=triton sm_120(RTX 50系)では flashinfer が hybrid-GDN パスで assert する
--kv-cache-dtype int2 上述のとおり、無いと262Kが入らない
MAMBA_SCHED=extra_buffer radix cacheを有効にしたままoverlapスケジューラを維持できる唯一の方法。no_buffer だと "Disabling overlap schedule" され集約スループットが14%低下
MEM=0.88 / MAXMAMBA=64 これ以上上げてはいけない。MEM=0.92/MAXMAMBA=80 だとプールが埋まるにつれGPUが11,791MiBに達し、VRAM崖(約11,830MiB)を越えてprefillが3,056→16 tok/s(190分の1)に崩壊した。0.88/64ならプールはむしろ大きく(350,257 > 303,751)、余裕が約1.1GB残る。代償は並列16→12だけ

--enable-mixed-chunk は使えません。INT2カーネルと併用すると CUDA error: invalid configuration argument でクラッシュします。

VRAM崖はWSLだと確保エラーとして表面化しません。GPU使用率100%・消費電力40〜50Wのまま「なんか遅い」に見えるだけです。数字に出ないので、崖の手前で止める設定を最初から選ぶしかありません。


第6部: Qwen3.8-27B(dense版)でも同じか

ほぼ同じです。同じ環境でdense版の Qwen3.8-27B-Escha-W2 も OSCAR INT2 KV 付きで動かしていて(196,608トークン、同じ12GB×2)、ステップ3の統合は同じ10ファイルに対する同じ作業でした。対応表はそのまま使えます。

違うのはこの表のところです。

MoE(本記事) dense Qwen3.8-27B
wheel escha-1.0.2+qwen3moe escha-1.2.1+qwen3dense
層数 / head_dim 40 / 256 64 / 256
full_attention層 10(4層に1層) 16(4層に1層)
TP2時の1GPUあたりヘッド 8 q / 1 kv → kv_group_num 8 12 q / 2 kv → kv_group_num 6
本番の回転行列 較正なし(Hadamard) 較正あり(層ごとに異なる)
KV量子化グループ --kv-cache-quant-group-size 32 64
expert並列 --ep-size 2 が必須 該当なし(denseなので)
--enable-mixed-chunk INT2カーネルと衝突してクラッシュ 使える(実際に使用中)
コンテキスト 262,144 196,608

着手前に知っておくべき差が3つあります。

1. 公開しているパッチはMoEのwheel向けです。 dense側のツリーに --dry-run で当ててみたところ、6本中3本が失敗しました(triton_backend、schedule_policy、scheduler の2つ目のhunk)。decode_attention も fuzz 付きでしか当たりません。変更内容はどちらでも正しいが、コンテキスト行が合わないだけなので、無理に当てずに手で移植して測り直してください。

2. _safe_block_h のバグはdense版でこそ効きます。 TP2時のdense版は kv_group_num = 6。BLOCK_H 以上でもその倍数でもない、まさにヘッドタイルがKVヘッドをまたぐ形です。MoE側は 8(2冪)なので構成上安全でした。同じコードでもモデルによって踏むか踏まないかが変わる典型例です。

3. 回転行列の選択をコピーしないでください。 MoEのチェックポイントでは較正なしが勝ちましたが、dense版の本番は較正ありで回しています。どちらが良いかはチェックポイントの性質なので、まずツールでHadamard版を作って動かし、較正に投資するのは「識別力のあるテスト」が効果を示してからにしてください。

起動フラグの差は、--ep-size 無しの --tp-size 2、--kv-cache-quant-group-size 64、--triton-attention-num-kv-splits 64(既定の8が動的分割数の上限を切っていて、長文脈デコードで約8%損をします)、--mamba-ssm-dtype float16、そして CTXLEN=196608 MEM=0.87 CHUNK=6144 MAXREQ=12 MAXMAMBA=40 です。

まとめ

  • 配布wheelに他人のPRを載せる作業は、パッチ適用ではなくマージです。ツリーが違う前提で、ファイル単位に「コピー/手マージ/移植しない」を判断する
  • OSCARの回転行列は較正不要のHadamardで十分。論文の劣化報告と逆の結果が出ることがあるので必ず測る。needleでは測れない
  • 匿名化した .patch は -p0 で当たらない。patch はスキップして正常終了するので気付けない
  • 「動いた」の確認は起動ログのKVサイズと import で
  • dense版(Qwen3.8-27B)でも同じ10ファイル・同じ作業。ただしパッチのコンテキスト行は合わず(6本中3本が失敗)、kv_group_num が6になるので _safe_block_h はdense版でこそ必要。回転行列の較正あり/なしの判断もモデルごとに測り直す

対応表・パッチ6本・回転行列生成スクリプトは以下にまとめてあります(Apache-2.0。eschaランタイムとモデル重みは再配布していません)。

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?