Skip to content

3bit EXL3 エキスパートのプレフィルを融合カーネルに通す(#22 の A1) - #26

Open
GarupanOjisan wants to merge 2 commits into
sync-upstream-1.1.5from
glm-exl3-3bit-fused-prompt
Open

GarupanOjisan wants to merge 2 commits into
sync-upstream-1.1.5from
glm-exl3-3bit-fused-prompt

Conversation

@GarupanOjisan

@GarupanOjisan GarupanOjisan commented Oct 26, 2026 •

Copy link
Copy Markdown
Owner

#22 の A1 を実装した。3bit EXL3 のルーティッドエキスパートの層も、プレフィルで融合カーネル(rotate、gate/up、down の 3 つ)を通るようにした。
#22 の候補のうち、この PR で扱うのは A1 だけである。

PR #23(upstream v1.1.5 の取り込み)の上に積んでいる。v1.1.5 が zig/src/families/glm/cuda_exl3.zig と zig/exl3_prompt_host.zig を変えているためである。#23 を先にマージしてほしい。
なお、#22 の B1(GLM の KV キャッシュのページング)は v1.1.5 に入っているので、#23 で片付く。

Refs #22

変えたこと

これまで融合経路に入れたのは 4bit の層だけだった。条件が uniform_k2 == 8 だったためである。
3bit の層(uniform_k2 = 6)は grouped 経路に回っていた。grouped 経路は 6 つのカーネルを別々に起動し、入力の Hadamard 回転を 1 行につき 8 回(選ばれたエキスパートの数)繰り返す。

融合カーネルで 4bit に固定されていたのは、タイルの語数(32 語)と、タイルを 16 バイトずつ 8 回で読む部分だけだった。
そこで、zig/kernels/cuda/exl3_prompt.cuh の decode_tile_p、tmain、tgateup を trellis の幅 K2 でテンプレート化した。
3bit のタイルは 24 語なので、16 バイトずつ 6 回で読む。共有メモリのタイルの枠は 32 語のままにしている。
タイルの並び((kタイル × 列タイル数 + 列タイル) × 語数)は grouped カーネルと同じである。

  • zig/kernels/cuda/exl3_experts.cu:tf_exl3_prompt_gu3 と tf_exl3_prompt_down3(K2 = 6)を足した。4bit の tf_exl3_prompt_gu と tf_exl3_prompt_down は、名前と引数を変えずに同じテンプレートから作っている。
  • zig/src/cuda/exl3_prompt.zig:eligible が uniform_k2 の 6 も受け付ける。run は 3bit なら新しい 2 つのカーネルを起動する。
  • 実際の GLM-5.3-Flash の層は、すでに融合経路を通っている 4bit 版と同じ形(dims % 256 == 0、width % 128 == 0)なので、ほかの条件は変わらない。

出力について

#22 には「計算は同じなので出力は変わらない見込み」と書いたが、これは正しくなかった。
融合経路と grouped 経路では、K 方向の足し合わせの順番が違う(grouped 経路は split-K の部分和を後で足す)。
そのため 3bit の層の出力は fp32 の丸め誤差の範囲で変わり、greedy の出力トークンが今の配信と一致するとは限らない。
精度は、下のテストで float64 の参照に対する許容誤差の内側にあることを確認した。

検証

spark1(DGX Spark、sm_121)で、nvcr.io/nvidia/pytorch:26.07-py3 の中で Zig 0.17.0 を使った。

  • zig fmt --check build.zig zig:成功
  • zig build -Doptimize=ReleaseFast -Dnvcc=/usr/local/cuda/bin/nvcc install native test:成功
  • zig build test-kernel-contract:成功
  • tf-cuda-test glm-exl3(GPU):成功。このテストを 4bit と 3bit の両方で回すようにした。3bit の融合プレフィルは 1、40、200、1030 行で、float64 の参照との差が許容誤差の 0.144 倍以内だった(4bit は 0.186 倍以内)。1024 行の区切りの前後の行が、その行だけで計算したときとビット単位で一致することも確認した。
  • 4bit のカーネルへの影響:変更前(20bc843d)と変更後の exl3_experts.cu を、ビルドと同じフラグ(-O3 --expt-relaxed-constexpr)で sm_121 向けにコンパイルして比べた(nvcc 13.3)。tf_exl3_prompt_gu は 128 レジスタ、tf_exl3_prompt_down は 127 レジスタで、共有メモリも含めて変更前と同じだった。新しい gu3 と down3 はどちらも 128 レジスタで、共有メモリは 4bit 版と同じ大きさである。
  • zig/exl3_prompt_host.zig:3bit のプレフィルが gu3 と down3 を起動し、grouped カーネルを起動しないこと、デコードは grouped のままであることを確かめるテストを足した。uniform_k2 = 6 を対象外としていた既存のテストは、対象に変えた。

まだ確認していないこと

  • 速度:2 台の GB10 で、3bit の層のプレフィル時間が実際に短くなるかは測っていない。配信を止める必要がある。2 台の GB10 で GLM-5.3-Flash-EXL3(3bit)を速くする取り込み候補の整理(3bit の融合プレフィル、KV のページング、FP4 プレフィル) #22 にある「4bit と 3bit で 1 区切りあたりのエキスパートの時間を比べる」計測で確かめる。
  • 実機の出力:GLM-5.3-Flash EXL3 3bit、TP=2 の greedy 出力が、今の配信と同程度の品質であることは見ていない。
  • 契約レジストリ(zig/src/core/contract/registry.zig):新しい 2 つのカーネルのレジスタ数と共有メモリの記録は足していない。記録には sm_80、86、89、120、121 の fatbin で tf-cuda-test kernel-contract-limits を回す必要がある。登録していないカーネルは検査の対象にならないので、test-kernel-contract は通る。
    なお、nvcc 13.3 では 4bit の gu が sm_121 で 128 レジスタになり、記録(127)と合わない。これはこの PR の前から同じで、この PR では記録を直していない。

🤖 Generated with Claude Code

GarupanOjisan and others added 2 commits October 26, 2026 11:29
融合プレフィル経路(rotate、gate/up、down の 3 カーネル)を trellis の幅でテンプレート化し、
3bit(uniform_k2 6、24 語のタイル)用の tf_exl3_prompt_gu3 と tf_exl3_prompt_down3 を足した。
これまで 3bit の層は grouped 経路の 6 カーネルに回り、入力の Hadamard 回転を 1 行につき 8 回繰り返していた。

Refs #22

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant